PyTorch Training CNN on MNIST Dataset

PyTorch CNN Commands Cheat Sheet

1. Set hyperparameters—these are safe to start with.

2. Specify local drive folders to store the MNIST dataset, and a location for the trained data.

3. The MNIST dataset comes built into PyTorch, accessible via torchvision.datasets.MNIST. Transform the dataset using the transforms.Compose() function, as follows:

4. Create the train_dataset and test_dataset objects, providing the following arguments:

• root—specifies the folder where the train.pt and test.pt data files exist

• train—specifies whether to use the train.pt or test.pt data file

• transform—passes the transform object, created earlier

• download—specifies that the MNIST data should be downloaded from source if needed

5. Load the train and test datasets into the data loader. A data loader can be used as an iterator – to extract the data, just use a standard Python iterator such as enumerate. The DataLoader function takes three arguments:

• The data set you wish to load

• The batch size

• Whether you wish to randomly shuffle the data

6. Create the CNN model by initializing the nn.Module class. This is a PyTorch class which has everything you need to build a neural network. It also provides recursive operations, ways of parallelizing work and moving it to a GPU or back to a CPU, and more. We’ll create the following neural layers:

• layer1—using the nn.Sequential object, we create a compound layer that includes a 2D convolutional layer, a ReLu activation function and a 2D MaxPool layer.

  • The Conv2d method lets you define, in this order: number of input channels (1 because MNIST images are grayscale), number of output channels, the size of the convolutional filter (you can supply a tuple for different shapes), stride and padding.

  • Stride and padding for the convolutional layer are defined using this equation: Wout=(Win–F+2P)S+1 where Wout=width of output, Win = width of input, F=filter, S=stride, P=padding. The same formula can be used for height, since the images are symmetric. Because we want the output size to be the same as the input (2), we set stride to 1 and padding to 2 → the output width and height = 2.

  • For the pooling layer, we set stride to 2 and padding to zero, to down-sample and reduce images by a factor of 2.

  • Thus, the output from layer1 will be 32 channels of 14 x 14 pixel images.

• layer2—like 1, except input channels are 32 because it received the output of the first layer, and output 64 channels.

• drop_out layer to avoid over-fitting in the model.

• fc1 and fc2—Two fully connected layers, created using the nn.Linear method. The first with 7 x 7 x 64 nodes, and the second with 1000 nodes. The second argument of nn.Linear specifies the number of nodes in the next layer.

7. Define the forward pass. To customize forward pass functionality, call this function “forward”, to override the base forward function in nn.Module. The second argument x is one batch of data, which is fed into the first neural layer (layer1), then to the next layer, and so on. After layer2, we reshape the data, flattening it from 7 x 7 x 64 into 3164 x 1.

8. Define training parameters by creating a ConvNet object and defining:

• criterion—the loss function. We use the PyTorch CrossEntropyLoss function which combines a SoftMax and cross-entropy loss function.

• optimizer—we use the Adam optimizer, passing all the parameters from the CNN model we defined earlier, and a learning rate.

9. Train the model, by running two loops:

• Loop over the number of epochs.

• Within this loop, iterate over train_loader using enumerate, and do the following:

  • Perform a forward pass, by passing a batch of normalized MNIST images from train_loader to the model object you defined earlier. There is no need to explicitly run the forward function, PyTorch does this automatically when it executes a model.

  • Pass the outputs true image labels to the loss function.

  • Append the loss to a list, which you can use later to plot training progress.

  • In preparation for backpropagation, set gradients to zero by calling zero_grad() on the optimizer.

  • Perform backpropagation using the backward() method of the loss object. Gradients are calculated.

  • Call optimizer.step() to perform Adam optimizer training.

10. Track accuracy (within the same loop) by:

• Running the torch.max() function, which returns the index of the maximum value in a tensor.

• In the first argument of the max function, pass a tensor of outputs from the model, which should be of size (batch_size, 10). For each sample in the batch, this will return the maximum value over the 10 output nodes, each representing one of the digits 0-9. For example, output 2 corresponds to digit “2”. The node with the highest output value will be predicted by the model.

• In the second argument of the max() function, pass 1. This instructs the max function to examine the output node axis (axis=0 corresponds to the batch_size dimension).

• max() returns a list of prediction integers from the model

• Compare predictions with the true labels (predicted == labels) and sum them to get the number of correct predictions. Divide by the batch_size to obtain the accuracy.

• Print progress after every 100 iterations.

11. To test the model, do the following:

• Run model.eval()—this runs the model while disabling drop-out or batch normalization layers.

• torch.no_grad() disables autograd functionality in the model, this is PyTorch’s mechanism for performing backpropagation and calculating gradients, which is not needed in model testing.

• The remaining code is the same as in the accuracy calculation above, except you are iterating through test_loader and not train_loader.

• Output the prediction to the console, save model results using torch.save(). This enables graphing the results using a plotting library.