Chapter 7: Building and Training a simple Classification Model

Implementing Simple Classifier

face Luca Grillotti

Before Implementing anything…

There are a few questions we need to address before starting to implement our MNISTClassifier module.

We have already answered most of these questions last chapter, but let’s recap quickly:

What kind of model should we use?

For the moment, we just consider that our model performs a single linear transformation to the image, without anything else (no activation function, no additional linear transformations…)

No worries, we will add more layers very soon 😉

What is the input to our model?

Our model should take as input a batch of images of size: torch.Size([N_batch, 1, 28, 28]), where N_batch represents the number of images in the batch.

What is the output of our model?

Output Characteristics

For each image we expect our model to output the score s_i of belonging to each class \mathcal{C}_i. Where:

  • \mathcal{C}_i is the class of digits i
  • 0 \leq i \leq 9.

In other words, for each image, our model should output a vector of scores: [s_0, s_1, \cdots, s_9].

And if we feed the model with an image of a 5, we expect our model to output a vector of 10 scores: [s_0, s_1, \cdots, s_9] where s_5 is significantly higher than the other scores.

Output Shape

As explained above, for each image, we expect our model to output a 10-dimensional vector.

So, when we give as input to the model N_batch images, we expect our model to return N_batch 10-dimensional vectors.

Thus, if the input is of shape torch.Size([N_batch, 1, 28, 28]), the output should be of shape torch.Size([N_batch, 10]).

Which loss function should we consider?

Converting scores into probabilities:

Given an image I, our model should output a 10-dimensional vector of scores [s_0, s_1, \cdots, s_9].

Those scores s_i can take any value in the set of real numbers. But we would like to convert them into a vector of probabilities: [p(I=0), p(I=1), \cdots, p(I=9)], where:

  • p(I=0) is the predicted probability that the label of I is 0.
  • p(I=1) is the predicted probability that the label of I is 1.
  • and so on…

One common approach consists of using the softmax function to estimate those probabilities:

p(I=i) = \dfrac{\exp(s_i)}{\sum_{j=0}^{9} \exp(s_j)}.

Comparing predicted probabilities with ground-truth labels

Ideally, if our model was perfect, given an image of a 0, our model should predict the vector of probabilities: [1, 0, 0, ...].

So, given an image of a 0, we intend to minimise the distance between the predicted probabilities [p(I=0), p(I=1), \cdots, p(I=9)] and the ground truth label of the image [1, 0, 0, ...].

Cross-Entropy Loss Function

In PyTorch, the loss torch.nn.CrossEntropyLoss performs the two operations we just mentioned:

  1. Converting the scores returned by the model [s_0, s_1, \cdots, s_9] into a vector of probabilities [p(I=0), p(I=1), \cdots, p(I=9)].
  2. Calculating the distance between the estimated probabilities [p(I=0), p(I=1), \cdots, p(I=9)] with

In the end, by minimising the cross-entropy loss, we improve the predictions of our model.

And in PyTorch?

Remember the structure of our good old PyTorch Module?

import torch

class MNISTClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        ...

    def forward(self, tensor_images):
        """
        Args:
            tensor_images: tensor of shape (N_batch, 1, 28, 28)
        """
        ...

We now have all the tools to fill this module, and adapt it to our problem! :D

The __init__ method

As we said previously, for the moment, we just consider our model is made of a single linear layer.

    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(in_features=<in_features>, out_features=<out_features>)

What should we put for <in_features> and <out_features>?

<out_features> is easy to determine, we said we wanted our model to output a 10-dimensional vector for each image. So we take <out_features>=10.

<in_features> corresponds to the total number of dimensions of each image. As each image is made of 1 channel with 28 * 28 pixels: <in_features>=1 * 28 * 28.

The forward method

    def forward(self, tensor_images):
        """
        Args:
            tensor_images: tensor of shape (N_batch, 1, 28, 28)
        """
        tensor_images = tensor_images.view(-1, 1 * 28 * 28)
        outcome_scores = self.linear(tensor_images)
        return outcome_scores

In Summary:

class MNISTClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(in_features=1 * 28 * 28, out_features=10)

    def forward(self, tensor_images):
        """
        Args:
            tensor_images: tensor of shape (N_batch, 1, 28, 28)
        """
        tensor_images = tensor_images.view(-1, 1 * 28 * 28)
        outcome_scores = self.linear(tensor_images)
        return outcome_scores

Let’s test!

It is always important to test our model by providing random tensors as input. This way, we ensure that we did not do any mistake.

import torch
mnist_classifier = MNISTClassifier()
n_batch = 42
random_tensor = torch.randn(size=(n_batch, 1, 28, 28))
print(mnist_classifier(random_tensor).size()) 

and we get:

torch.Size([42, 10]) 

so, the result is of size (n_batch, 10), as expected! \o/