MNIST: the Hello World Example in Image Recognition

The crucial difference between learning and pure optimization is that the dataset of interest \(\mathcal{D}\) is unknown in a learning problem. Indeed, we cannot know the data encountered in applications and their labels before deploying our model. In most cases, we only have access to a training dataset \(\mathcal{D}_{\text{train}}\) and do our work with it. As \(\mathcal{D}_{\text{train}}\) may not fit the true data density, it is often necessary to preserve a part of it to avoid overfitting, which is called the validation dataset. Nevertheless, for the sake of simplicity we do not use this technique in this post. As the MNIST dataset provides the labels of both the training set and test set, we use the test accuracy to evaluate our model performance directly. However, it should be keep in mind that if there is no label data in the test dataset then one has to split the training dataset to construct a validation dataset manually.

We usually do a simple preprocess when loading the data, e.g., a normalization to scale the data to have mean 0 and std 1. The original mean and std of MNIST training set can be calculated easily by the following python statement.

x = torch.cat([train_data[i][0] for i in range(len(train_data))], dim=0)
print(x.mean().item(), x.std().item())

0.13066047430038452 0.30810782313346863

As we will use the idea of SGD to optimize the objective function, it is convenient to create a data loader to iteratively select a batch of data points.

train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST(
        "./data/",
        train=True,
        download=True,
        transform=torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize((0.1307,), (0.3081,)),
            ]
        ),
    ),
    batch_size=cfg["batch_size_train"],
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST(
        "./data/",
        train=False,
        download=True,
        transform=torchvision.transforms.Compose(
            [
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize((0.1307,), (0.3081,)),
            ]
        ),
    ),
    batch_size=cfg["batch_size_test"],
    shuffle=True,
)

Doing so allows us to use for loop to iterate the training dataset conveniently by writing

for batch_x, batch_y in train_loader:
    # do SGD with this batch of data
    pass

What happens behind this is:

1. the order of samples in the dataset is shuffled before the iteration;
2. each sample get preprocessed by normalizing with mean 0.1307 and std 0.3081;
3. in each step, a fixed number of samples are drown from the dataset and are stacked into batch_x and batch_y.

An epoch means a whole for loop and a batch means a step in the for loop.

From the computational view, a neural network \(f(x;\theta)\) is a nested function \[ f(\cdot; \theta) = f_{N-1} \circ f_{N-2} \circ \cdots f_0,\] where each layer \(f_t\) is a parameterized function with parameter \(\theta_t\). Then the parameter \(\theta\) of the neural network \(f(x;\theta)\) is actually the collection \(\{\theta_t,\ t=0,1,\ldots, N-1\}\).

In PyTorch, a convolutional layer is a function which accepts a 4D tensor \(x[\alpha,i,j,k]\) and outputs another 4D tensor \(y[\alpha, i', j', k']\). The parameter of th layer consists of a bias tensor \(b[i', j', k']\) and a weight tensor \(w[i', i, j'', k'']\). In particular, \[ y[\alpha, i'] = b[i'] + \sum_{i} w[i', i] \star x[\alpha, i],\] where \(\star\) is the correlation operator between matrices; see also the documention of torch.nn.Conv2d for more details. This post also gives a review on how torch.nn.Conv2d works.

Compared to convolutional layers, fully connected layers, i.e., linear layers in PyTorch is rather simple. They are just affine transformations. In the most simple case, a linear layer in PyTorch accpets a 2D tensor \(x[\alpha, i]\) and outputs another 2D tensor \(y[\alpha, i']\). The parameter of the layer consists of a bias tensor \(b[i']\) and a weight tensor \(w[i', i]\). In particular, \[ y[\alpha] = b + w \circ x[\alpha], \] where \(\circ\) is the matrix-vector product. In the general case, the index \(\alpha\) might be multiple indices and the input \(x\) and \(y\) become high order tensors; see also the documentation of torch.nn.Linear.

We use CNN as the basic model of our classifier[3]. In particular, the model consists of two 2D convolutional layers followed by two fully-connected layers. After each convolutional layer, there is a maximum pooling operation. In addition, the activation function is called between any two layers. We choose ReLU as our activation function.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
        self.maxpool = nn.MaxPool2d(kernel_size=2)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.maxpool(self.conv1(x)))
        x = self.relu(self.maxpool(self.conv2(x)))
        x = x.view(-1, 320)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x


clf = Net()

As we can see, the output of our model \(f(x; \theta)\) is a vector with 10 components. But how to predict the label of \(x\) and evaluate its performance? The de facto standard way is interpreting the components as the logit of the class. For example, in MNIST there are 10 classes, i.e., 10 labels in total. If the model returns \((t_0, t_1, \ldots, t_{9})\) for an image, then we say the model predicts that the probability distribution of the label \(y\) \[ \mathbb{P}(y = i) = \frac{e^{t_i}}{\sum_i e^{t_i}},\quad i=0,1,\ldots,9. \] Let \(p_i = \mathbb{P}(y=i)\) be the predicted distribution. We evaluate its performance by the relative entropy of \(p\) with respect to the true distribution \(q\), i.e., \[\ell(f(x;\theta), y) = -\sum_{i}q_i\log p_i,\] where the true distribution \(q_i\) is, of course, a deterministic distribution

Fortunately, PyTorch provides a convenient class CrossEntropyLoss to carry out above calculations given the predicted logits \(f(x;\theta)\) and the true label \(y\).

loss_func = nn.CrossEntropyLoss()

Given a pair of data \((x^{(i)}, y^{(i)})\), assume the loss of this pair is \(\phi(\tilde{y}^{(i)}, y^{(i)})\) where \(\tilde{y}^{(i)}=f(x^{(i)}; \theta)\), we want to compute the gradient of it w.r.t. \(\theta\) in order to perform gradient descent. For example, the mean-square error corresponds to \(\phi(\tilde{y}, y)=\|\tilde{y} - y\|^2\). Below is a brief summary of this post on back propagation.

Let us slightly overload the notation to denote by \(f_t(\cdot) = f(t, \cdot, \theta_t)\). Introduce the Hamiltonian \[ H(t, x, u, p) = p^\intercal f(t, x, u).\] In the calculation of the gradient, the forward phase is first executed to obtain the state variables

Clearly, this is identical to \(x_{k+1} = f_k\circ f_{k-1}\circ \cdots f_0(x^{(i)})\). Hence, \(x_{N} = f(x^{(i)}; \theta) = \tilde{y}\) and the loss is \(\phi(x_N, y)\). The backward phase is then executed to obtain the costate variables

Clearly, this is identical to \(p_{t} = (\nabla_x f_t(x_t; \theta_t))^\intercal p_{t+1}\). Here \(\nabla_x f_t\) is a Jacobian matrix and \((\nabla_x f_t)^\intercal p_{t+1}\) is often computed efficiently via Jacobian-vector product.

Finally, it is not hard to show by induction that the gradient of loss is \[ \frac{\partial}{\partial \theta_t}\phi(f(x^{(i)}; \theta), y^{(i)}) = \nabla_u H(t, x_t, u, p_{k+1})\big\vert_{u=\theta_t},\quad t=0, 1, \ldots, N-1. \]

We use the stochastic gradient descent algorithm to find the best net parameters \(\theta\). There is no need to compute the gradient by ourselves as PyTorch has implemented the back propagation algorithm internally and provides various optimizers in torch.optim package. We choose the simple SGD here.

optimizer = torch.optim.SGD(
    clf.parameters(), lr=cfg["learning_rate"], momentum=cfg["momentum"]
)