Marie-Hélène Burle
This classic PyTorch tutorial goes over the entire workflow to create and train a simple image classifier.
Let’s go over it step by step.
CIFAR-10 from the Canadian Institute for Advanced Research is a classic dataset of 60,000 color images falling into 10 classes (6,000 images in each class):
The images are of size 32x32 pixels (tiny!), which makes it very lightweight, quick to load and easy to play with.
A DataLoader is an iterable feeding data to a model at each iteration. The data loader transforms the data to the proper format, sets the batch size, whether the data is shuffled or not, and how the I/O is parallelized. You can create DataLoaders with the torch.utils.data.DataLoader class.
Let’s create 2 DataLoaders: one for the train set and one for the test set.
--------------------------------------------------------------------------- ModuleNotFoundError Traceback (most recent call last) Cell In[1], line 1 ----> 1 import torch 2 import torchvision 3 import torchvision.transforms as transforms ModuleNotFoundError: No module named 'torch'
The CIFAR-10 images in the TorchVision library are Image objects (from the PIL.Image module of the pillow package).
We need to normalize them and turn them into tensors:
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[2], line 1 ----> 1 transform = transforms.Compose( 2 [transforms.ToTensor(), 3 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) NameError: name 'transforms' is not defined
Remember that the data move forward through the network (forward pass), outputting some estimates which are used to calculate some loss (or error) value. Then we get gradients through automatic differentiation and the model parameters are adjusted a little through gradient descent.
You do not have to have the entire training set go through this process each time: you can use batches.
The batch size is the number of items from the data that are processed before the model is updated. There is no hard rule to set good batch sizes and sizes tend to be picked through trial and error.
Here are some rules to choose a batch size:
Let’s set the batch size to 4:
trainset = torchvision.datasets.CIFAR10(root='./data',
train=True,
download=True,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data',
train=False,
download=True,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=2)We will also need the classes:
Though not necessary, it can be useful to have a look at the data:
# Load the packages for this
import matplotlib.pyplot as plt
import numpy as np
# Define a function to display an image
def show_img(img):
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.figure(figsize=(9.5, 5))
plt.axis('off')
plt.imshow(np.transpose(npimg, (1, 2, 0)))
plt.show()
# Get a batch of random training images
dataiter = iter(trainloader)
images, labels = next(dataiter)
# Display the images
show_img(torchvision.utils.make_grid(images))
# Print the labels
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))--------------------------------------------------------------------------- ModuleNotFoundError Traceback (most recent call last) Cell In[6], line 2 1 # Load the packages for this ----> 2 import matplotlib.pyplot as plt 3 import numpy as np 4 5 # Define a function to display an image ModuleNotFoundError: No module named 'matplotlib'
First, we need to define the architecture of the network. There are many types of architectures. For images, CNN are well suited.
In Python, you can define a subclass of an existing class with:
The subclass is derived from the base class and inherits its properties. PyTorch contains the class torch.nn.Module which is used as the base class when defining a neural network.
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
# Define the architecture of the network
def __init__(self):
super().__init__()
# 3 input image channel (3 colour channels)
# 6 output channels,
# 5x5 square convolution kernel
self.conv1 = nn.Conv2d(3, 6, 5)
# Max pooling over a (2, 2) window
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
# 5*5 from image dimension
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
# 10 is the size of the output layer
# since there are 10 classes
self.fc3 = nn.Linear(84, 10)
# Set the flow of data through the network for the forward pass
# x represents the data
def forward(self, x):
# F.relu is the rectified-linear activation function
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
# flatten all dimensions except the batch dimension
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x--------------------------------------------------------------------------- ModuleNotFoundError Traceback (most recent call last) Cell In[7], line 1 ----> 1 import torch.nn as nn 2 import torch.nn.functional as F 3 4 class Net(nn.Module): ModuleNotFoundError: No module named 'torch'
Let’s create an instance of Net and print its structure:
We need to choose a loss function that will be used to calculate the gradients through backpropagation as well as an optimizer to do the gradient descent.
SGD with momentum has proved a very efficient optimizing technique and is widely used.
--------------------------------------------------------------------------- ModuleNotFoundError Traceback (most recent call last) Cell In[9], line 1 ----> 1 import torch.optim as optim 2 3 criterion = nn.CrossEntropyLoss() 4 optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) ModuleNotFoundError: No module named 'torch'
We can now train the model:
for epoch in range(2): # loop over the dataset twice
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 2000 == 1999: # print every 2000 mini-batches
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
print('Finished Training')--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[10], line 5 1 for epoch in range(2): # loop over the dataset twice 2 3 running_loss = 0.0 4 ----> 5 for i, data in enumerate(trainloader, 0): 6 # get the inputs; data is a list of [inputs, labels] 7 inputs, labels = data 8 NameError: name 'trainloader' is not defined
Let’s now test our model on one batch of testing data.
First, let’s get a batch of random testing data:
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[11], line 1 ----> 1 dataiter = iter(testloader) 2 images, labels = next(dataiter) NameError: name 'testloader' is not defined
Let’s display them and print their true labels:
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[12], line 1 ----> 1 show_img(torchvision.utils.make_grid(images)) 2 print('Real: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4))) NameError: name 'show_img' is not defined
Now, let’s run the same batch of testing images through our model:
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[13], line 1 ----> 1 outputs = net(images) NameError: name 'net' is not defined
Let’s get the best predictions for these:
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[14], line 1 ----> 1 _, predicted = torch.max(outputs, 1) 2 3 print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}' 4 for j in range(4))) NameError: name 'torch' is not defined
This was fun, but of course, with a sample of one, we can’t say anything about how good our model is. We need to test it on many more images from the test set.
Let’s use the entire test set:
correct = 0
total = 0
# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
for data in testloader:
images, labels = data
# calculate outputs by running images through the network
outputs = net(images)
# the class with the highest energy is what we choose as prediction
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[15], line 4 1 correct = 0 2 total = 0 3 # since we're not training, we don't need to calculate the gradients for our outputs ----> 4 with torch.no_grad(): 5 for data in testloader: 6 images, labels = data 7 # calculate outputs by running images through the network NameError: name 'torch' is not defined
We could see whether the model seems to perform better for some classes than others:
# prepare to count predictions for each class
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}
# again no gradients needed
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predictions = torch.max(outputs, 1)
# collect the correct predictions for each class
for label, prediction in zip(labels, predictions):
if label == prediction:
correct_pred[classes[label]] += 1
total_pred[classes[label]] += 1
# print accuracy for each class
for classname, correct_count in correct_pred.items():
accuracy = 100 * float(correct_count) / total_pred[classname]
print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[16], line 6 2 correct_pred = {classname: 0 for classname in classes} 3 total_pred = {classname: 0 for classname in classes} 4 5 # again no gradients needed ----> 6 with torch.no_grad(): 7 for data in testloader: 8 images, labels = data 9 outputs = net(images) NameError: name 'torch' is not defined