A Beginner’s Tutorial on Building an AI Image Classifier using PyTorch | by Alexander Wu | Towards Data Science

This is a step-by-step guide to build an image classifier. The AI model will be able to learn to label images. I use Python and Pytorch.
When we write a program, it is a huge hassle manually coding every small action we perform. Sometimes, we want to use packages of code other people have already written. These packaged routines are called Libraries and can be added into our program by importing them and then referencing the library later in the code.
We usually import all the libraries at the beginning of the program.
Instead of typing the library’s long name whenever we reference it, we can shorten it to a name of our choice using “as”.
# Import Librariesimport torchimport torchvision.transforms as transformsimport torchvision.datasets as datasetsimport torchvision.models as modelsimport torch.nn as nnimport torch.optim as optimimport numpy as npfrom PIL import Imageimport numpy as npimport matplotlib.pyplot as plt
Next, we want to import the picture data our AI model will learn from.
But before that, we need to specify the alterations we want to perform on these pictures — since the same command that imports them also transforms the data.
These transforms are made using the torchvision.transforms library. The best way to understand the transforms is to read the documentation here. But I’ll give a brief of what each command is doing.
transforms.Compose lets us compose multiple transforms together so we can use more than one transformation.
transforms.Resize((255)) resizes the images so the shortest side has a length of 255 pixels. The other side is scaled to maintain the aspect ratio of the image.
transforms.CenterCrop(224) crops the center of the image so it is a 224 by 224 pixels square image.
We do these two steps so all the images going into our AI model have the same size (AI models can’t handle inputs with varying sizes)
transforms.ToTensor() converts our image into numbers. How does it do that?
It separates the three colors that every pixel of our picture is comprised of: red, green & blue. This essentially turns one image into three images (one tinted red, one green, one blue).
Then, it converts the pixels of each tinted image into the brightness of their color, from 0 to 255. These values are divided by 255, so they can be in a range of 0 to 1. Our image is now a Torch Tensor (a data structure that stores lots of numbers).
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) subtracts the mean from each value and then divides by the standard deviation.
We will be using a pre-trained model, so we need to use the means and standard deviations the Pytorch specifies. There are three values in the mean and standard deviation to match each RGB picture.
# Specify transforms using torchvision.transforms as transforms# librarytransformations = transforms.Compose([    transforms.Resize(255),    transforms.CenterCrop(224),    transforms.ToTensor(),    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
Finally, we can import our pictures into the program. We use the torchvision.datasets library.
Read about it here.
We specify two different data sets, one for the images that the AI learns from (the training set) and the other for the dataset we use to test the AI model (the validation set).
The datasets.ImageFolder() command expects our data to be organized in the following way: root/label/picture.png. In other words, the images should be sorted into folders. For example, all the pictures of bees should be in one folder, all the pictures of ants should be in another etc.
We give the command
the path to all the folders andthe transforms that we specified in the previous step.
the path to all the folders and
the transforms that we specified in the previous step.
# Load in each dataset and apply transformations using# the torchvision.datasets as datasets librarytrain_set = datasets.ImageFolder("root/label/train", transform = transformations)val_set = datasets.ImageFolder("root/label/valid", transform = transformations)
Then we want to put our imported images into a Dataloader.
Dataloader is able to spit out random samples of our data, so our model won’t have to deal with the entire dataset every time. This makes training more efficient.
We specify how many images we want at once as our batch_size (so 32 means we want to get 32 images at one time). We also want to shuffle our images so it gets inputted randomly into our AI model.
Read about the DataLoader here.
# Put into a Dataloader using torch librarytrain_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)val_loader = torch.utils.data.DataLoader(val_set, batch_size =32, shuffle=True)
AI models need to be trained on a lot of data to be effective. Since we don’t have that much data, we want to take a pre-trained model (a model that has been previously trained on many images) but tailor it to recognize our specific images. This process is called transfer learning.
Image recognition models have two parts:
the convolutional part andthe classifier part
the convolutional part and
the classifier part
We want to keep the pre-trained convolutional part but put in our own classifier. Here’s why:
The convolution/pooling section of our model is used to recognize the features inside an image. It first identifies edges, then using the edges it identifies shapes, and using the shapes it can identify objects.
But it takes A LOT of data to train this section — probably more than we have — so instead, we can use the default, pre-trained convolutional layers. These pre-trained convolutional layers were trained to identify these feature very well, regardless of what kind of image you have.
There are also pooling layers in between convolutional layers that distill an image to a smaller size so it can be easily inputted to our classifier.
The last part of the model is the classifier. The classifier takes all the information extracted from the photo in the convolution part, and uses it to identify the image. This is the part of the pre-trained model we want to replace and to train on our own images. This makes the model tailored to identify the images we give it.
We use the torchvision.models library to download a pre-trained model. There are many different models we can download, and more info can be found here. I chose a model called densenet161 and specified that we want it to be pre-trained by setting pretrained=True.
Then, we want to make sure we don’t train this model since it is already trained, and we only want to train the classifier we will put in next. We tell the model not to calculate the gradients of any parameter since that is only done for training.
# Get pretrained model using torchvision.models as models librarymodel = models.densenet161(pretrained=True)# Turn off training for their parametersfor param in model.parameters():    param.requires_grad = False
Now we want to replace the default classifier of the model with our own classifier. Classifiers are fully connected neural networks, so to do this, we must first build our own neural network.
A neural network is just a method to find complex patterns and connections between numbers of the input and the output. In this case, it takes the features of the image that were highlighted by the convolution section to determine how likely the image is a certain label.
The first thing we want to do is to determine the amount of numbers inputted to our neural network. This has to match the amount of numbers that is outputted from the section before (the convolutional section). Since we didn’t change the convolutional section at all, the amount of numbers inputted in our classifier should be the same as the model’s default classifier.
Next, we want to determine the number of outputs. This number should match how many types of images you have. The model will give you a list of percentages, each corresponding to how certain the picture is to that label. So if you have images of bees, ants, and flies, there are 3 labels. and there should be 3 numbers in the output layer each corresponding to the probability of the input being a bee, ant, or fly.
Once we have those details, we use the torch.nn library to create the classifier. Information can be found here.
nn.Sequential can help us group multiple modules together.
nn.Linear specifies the interaction between two layers. We give it 2 numbers, specifying the number of nodes in the two layer.
For example, in the first nn.Linear, the first layer is the input layer, and we can choose how many numbers we want in the second layer (I went with 1024).
nn.ReLU is an activation function for hidden layers. Activation functions helps the model learn complex relationships between the input and the output. We use ReLU on all layers except for the output.
We repeat this for as many hidden layers as you want, with as many nodes as you want in each layer.
nn.LogSoftmax is the activation function for the output. The softmax function turns the outputted numbers into percentages for each label, and the log function is applied to make it computationally faster. We must specify that the output layer is a column, so we set dimension equal to 1.
After creating our own classifier, we replace the model’s default one.
# Create new classifier for model using torch.nn as nn libraryclassifier_input = model.classifier.in_featuresnum_labels = #PUT IN THE NUMBER OF LABELS IN YOUR DATAclassifier = nn.Sequential(nn.Linear(classifier_input, 1024),                           nn.ReLU(),                           nn.Linear(1024, 512),                           nn.ReLU(),                           nn.Linear(512, num_labels),                           nn.LogSoftmax(dim=1))# Replace default classifier with new classifiermodel.classifier = classifier
Now the model is created! Next, we just need to train it.
Training a model on a GPU is a lot faster than a CPU. So to determine which device is available for you, we use Torch to check. If there is a compatible GPU, we set the variable to GPU, if not, it sticks with CPU. We then move our model to this device.
# Find the device available to use using torch librarydevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")# Move model to the device specified abovemodel.to(device)
While training, we need to determine how “off” our model is. To evaluate the amount of error our model has, we use nn.NLLLoss. This function takes in the output of our model, for which we used the nn.LogSoftmax function.
To train our model, we take our error and see how we can adjust the weights we multiplied our numbers by to get the smallest error. The method of calculating how we adjust our weights and applying it to our weights is called Adam. We use the torch.optim library to use this method and give it our parameters.
# Set the error function using torch.nn as nn librarycriterion = nn.NLLLoss()# Set the optimizer function using torch.optim as optim libraryoptimizer = optim.Adam(model.classifier.parameters())
Now we train. We want our model to go through the entire dataset multiple times, so we use a for loop. Every time it has gone over the entire set of images, it is called an epoch. In one epoch we want the model to go through both the training set and the validation set.
We first set the model to training mode and we use a for-loop to go through every image.
After moving the images and the labels to the appropriate device, we need to clear the adjustments of the weights by declaring optimizer.zero_grad().
We can then compute the output of our model given our images and how “off” our model is given its output and the correct answers.
Then we can find the adjustments we need to make to decreases this error by calling loss.backward() and use our optimizer to adjust the weights by calling optimizer.step().
As we train, we want to know how things are going, so we keep track of the total errors we calculated and print out the progress of the training.
We set our model to evaluation mode and use a for-loop to iterate over all the images in our set.
We repeat the steps we took for the training set to get the output of our model and how much our model is “off” from the real labels.
Our model used the LogSoftmax function in order to increase computation speed, but now we want the real percentages, not the log percentages. So we use torch.exp to reverse the log function. We then want to see which class the model guessed for our images. .topk gives us the top class that was guessed, and what percentage it guessed it at — we only care about the class so we can ignore the percentage.
To determine how many images it guessed right, we check which guessed classes are equal to the real classes. Then we can average over the entire batch to determine the accuracy of our model (how many images it guessed right divided by the total amount of images).
After going through both the training and validation set, we want to print the errors for both and the accuracy of the validation set.
epochs = 10for epoch in range(epochs):    train_loss = 0    val_loss = 0    accuracy = 0        # Training the model    model.train()    counter = 0    for inputs, labels in train_loader:        # Move to device        inputs, labels = inputs.to(device), labels.to(device)        # Clear optimizers        optimizer.zero_grad()        # Forward pass        output = model.forward(inputs)        # Loss        loss = criterion(output, labels)        # Calculate gradients (backpropogation)        loss.backward()        # Adjust parameters based on gradients        optimizer.step()        # Add the loss to the training set's rnning loss        train_loss += loss.item()*inputs.size(0)                # Print the progress of our training        counter += 1        print(counter, "/", len(train_loader))            # Evaluating the model    model.eval()    counter = 0    # Tell torch not to calculate gradients    with torch.no_grad():        for inputs, labels in val_loader:            # Move to device            inputs, labels = inputs.to(device), labels.to(device)            # Forward pass            output = model.forward(inputs)            # Calculate Loss            valloss = criterion(output, labels)            # Add loss to the validation set's running loss            val_loss += valloss.item()*inputs.size(0)                        # Since our model outputs a LogSoftmax, find the real             # percentages by reversing the log function            output = torch.exp(output)            # Get the top class of the output            top_p, top_class = output.topk(1, dim=1)            # See how many of the classes were correct?            equals = top_class == labels.view(*top_class.shape)            # Calculate the mean (get the accuracy for this batch)            # and add it to the running accuracy for this epoch            accuracy += torch.mean(equals.type(torch.FloatTensor)).item()                        # Print the progress of our evaluation            counter += 1            print(counter, "/", len(val_loader))        # Get the average loss for the entire epoch    train_loss = train_loss/len(train_loader.dataset)    valid_loss = val_loss/len(val_loader.dataset)    # Print out the information    print('Accuracy: ', accuracy/len(val_loader))    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(epoch, train_loss, valid_loss))
That was great! You’ve just built an AI image classifier. But, now we want to actually use it — we want to give it a random image and see which label it thinks it is.
First, we set the model for evaluation mode.
model.eval()
Then we create a function that can process the image so it can be inputted into our model.
We open the image, resize it by keeping the aspect ratio but making the shortest side only 255 px, and crop the center 224px by 224px. We then turn the picture into an array and make sure that the number of color channels is the first dimension instead of the last dimension by transposing the array. Next, we convert each value between 0 and 1 by dividing by 255. We then normalize the values by subtracting the mean and dividing by the standard deviation. Lastly, we convert the array into a Torch tensor and convert the values to float.
These steps are the same steps we specified in Step 2, but this time we must manually code the commands instead of relying on the transforms library.
# Process our imagedef process_image(image_path):    # Load Image    img = Image.open(image_path)        # Get the dimensions of the image    width, height = img.size        # Resize by keeping the aspect ratio, but changing the dimension    # so the shortest size is 255px    img = img.resize((255, int(255*(height/width))) if width < height else (int(255*(width/height)), 255))        # Get the dimensions of the new image size    width, height = img.size        # Set the coordinates to do a center crop of 224 x 224    left = (width - 224)/2    top = (height - 224)/2    right = (width + 224)/2    bottom = (height + 224)/2    img = img.crop((left, top, right, bottom))        # Turn image into numpy array    img = np.array(img)        # Make the color channel dimension first instead of last    img = img.transpose((2, 0, 1))        # Make all values between 0 and 1    img = img/255        # Normalize based on the preset mean and standard deviation    img[0] = (img[0] - 0.485)/0.229    img[1] = (img[1] - 0.456)/0.224    img[2] = (img[2] - 0.406)/0.225        # Add a fourth dimension to the beginning to indicate batch size    img = img[np.newaxis,:]        # Turn into a torch tensor    image = torch.from_numpy(img)    image = image.float()    return image
After processing the image, we can build a function to use our model to predict the label. We input the image into our model and obtain the output. We then reverse the log in the LogSoftmax function that we applied in the output layer and return the top class the model predicted and how certain it is of its guess.
# Using our model to predict the labeldef predict(image, model):    # Pass the image through our model    output = model.forward(image)        # Reverse the log function in our output    output = torch.exp(output)        # Get the top predicted class, and the output percentage for    # that class    probs, classes = output.topk(1, dim=1)    return probs.item(), classes.item()
Lastly, we want to display the image. We turn the image back into an array, and un-normalize it by multiplying by the standard deviation and adding back the mean. We then use the matplotlib.pyplot library to plot the picture.
# Show Imagedef show_image(image):    # Convert image to numpy    image = image.numpy()        # Un-normalize the image    image[0] = image[0] * 0.226 + 0.445        # Print the image    fig = plt.figure(figsize=(25, 4))    plt.imshow(np.transpose(image[0], (1, 2, 0)))
Now, we can use all these functions to print our model’s guess and how sure it was!
# Process Imageimage = process_image("root/image1234.jpg")# Give image to model to predict outputtop_prob, top_class = predict(image, model)# Show the imageshow_image(image)# Print the resultsprint("The model is ", top_prob*100, "% certain that the image has a predicted class of ", top_class  )