One of the standard PyTorch examples is to create a classification model for the MNIST dataset, using a convolutional neural network (CNN). If you’re new to PyTorch and you search the Internet for the MNIST CNN example, you’ll get the impression that it’s a simple problem. It’s not — MNIST CNN is extremely difficult.

Here’s my latest version. I don’t like using the built-in black box MNIST dataset from the torchvision library. That hides at least 50% of the key ideas, and in a real problem you don’t have a magic dataset available. Creating the training and test datasets from the binary source files is a major challenge in itself. I created a 1000-item training file and a 100-item test file. Each line is one 28×28 MNIST image of a handwritten digit. Each of the 28 x 28 = 784 pixel values is between 0 and 255. I structured the files so that the 784 pixel values come first, and the label (‘0’ through ‘9’) is the last value on each line. See https://jamesmccaffreyblog.com/2022/01/21/working-with-mnist-data/.

A typical MNIST digit

My network architecture uses two convolutional layers, each with 2×2 max pooling. After the convolutional layers, the network uses three linear layers.

Dealing with the sizes of the input and all these layers is another major challenge that can take hours or days to figure out, depending on your experience level. Briefly, I used a batch size of 10 and each input is 10 x 1 x 28 x 28 where the 1 comes from the fact that MNIST data is 1-channel grayscale rather than 3-channel color.

The really tricky size issue is the connection between the final convolution layer (after the second max pooling) and the first linear layer. In theory it’s possible to calculate this size but in practice it’s easier to put in a print() statement during development. For my architecture, the transition size is 64 * 4 * 4 = 1024.

My demo program does not use dropout. I only use dropout when I see overfitting and that doesn’t happen with my demo. Literally every other example I’ve seen on the Internet automatically uses dropout, slavishly following convention without really understanding the point.

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()  # pre Python 3.3 syntax

    self.conv1 = T.nn.Conv2d(1, 32, 5)  # chnl-in, out, krnl
    self.conv2 = T.nn.Conv2d(32, 64, 5)
    self.pool1 = T.nn.MaxPool2d(2, 2)   # kernel, stride
    self.pool2 = T.nn.MaxPool2d(2, 2)
    self.fc1 = T.nn.Linear(64 * 4 * 4, 512)  # empirically
    self.fc2 = T.nn.Linear(512, 256)
    self.fc3 = T.nn.Linear(256, 10)  # 10 classes
 
  def forward(self, x):
    # convolution phase            # x is Size([bs, 1, 28, 28])
    z = T.nn.functional.relu(self.conv1(x))  # [bs, 32, 24, 24]
    z = self.pool1(z)                        # [bs, 32, 12, 12]
    z = T.nn.functional.relu(self.conv2(z))  # [bs, 64, 8, 8]
    z = self.pool2(z)                        # [bs, 64, 4, 4]
   
    # neural network phase
    z = z.view(-1, 64 * 4 * 4)             # Size([bs, 1024])
    z = T.nn.functional.relu(self.fc1(z))  # Size([bs, 512])
    z = T.nn.functional.relu(self.fc2(z))  # Size([bs, 256])
    z = self.fc3(z)                        # Size([bs, 10])
    return z

Over the past two years or so, deep learning has become specialized: image recognition, tabular data, and natural language problems have become so complex that most of the guys I work with now mostly specialize in one of these three areas — the problems are just too complicated to be an expert in all three.

Good fun!

As recently as just a few years ago, recognizing even very simple images like 28×28 grayscale MNIST digits was a major challenge. But advances in CNN architecture and neural libraries like PyTorch, high accuracy can now be achieved on images like MNIST. But recognizing more complex images, such as these three abstract portraits, is still difficult for ML systems. Left: By artist Olga Kamennoy. Center: By artist Hans Jochem Bakker. Right: By artist Yossi Kottler.

Demo code. Replace “lt” (less-than”), “lte”, etc. with symbols — my lame blog editor often chokes on symbols.

# mnist_cnn.py
# PyTorch 1.10.0-CPU Anaconda3-2020.02  Python 3.7.6
# Windows 10 

# reads MNIST data from text file rather than using
# built-in black box Dataset from torchvision

import numpy as np
import torch as T

device = T.device("cpu")

# -----------------------------------------------------------

class MNIST_Dataset(T.utils.data.Dataset):
  # 784 tab-delim pixel values (0-255) then label (0-9)
  def __init__(self, src_file, n_rows=None):
    all_xy = np.loadtxt(src_file, max_rows=n_rows,
      usecols=range(785), delimiter="\t",
      skiprows=0, comments="#", dtype=np.float32)

    n = len(all_xy)
    tmp_x = all_xy[:, 0:784]  # all rows, cols [0,783]
    tmp_x /= 255
    tmp_x = tmp_x.reshape(-1, 1, 28, 28)  # len, chnl, 28x28
    tmp_y = all_xy[:, 784]    # 1-D required

    self.x_data = \
      T.tensor(tmp_x, dtype=T.float32).to(device)
    self.y_data = \
      T.tensor(tmp_y, dtype=T.int64).to(device) 

  def __len__(self):
    return len(self.x_data)

  def __getitem__(self, idx):
    lbl = self.y_data[idx]  # no use labels
    pixels = self.x_data[idx] 
    return (pixels, lbl)

# -----------------------------------------------------------

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()  # pre Python 3.3 syntax

    self.conv1 = T.nn.Conv2d(1, 32, 5)  # chnl-in, out, krnl
    self.conv2 = T.nn.Conv2d(32, 64, 5)
    self.pool1 = T.nn.MaxPool2d(2, 2)   # kernel, stride
    self.pool2 = T.nn.MaxPool2d(2, 2)
    self.fc1 = T.nn.Linear(64 * 4 * 4, 512)  # see text
    self.fc2 = T.nn.Linear(512, 256)
    self.fc3 = T.nn.Linear(256, 10)  # 10 classes
 
  def forward(self, x):
    # convolution phase            # x is Size([bs, 1, 28, 28])
    z = T.nn.functional.relu(self.conv1(x))  # [bs, 32, 24, 24]
    z = self.pool1(z)                        # [bs, 32, 12, 12]
    z = T.nn.functional.relu(self.conv2(z))  # [bs, 64, 8, 8]
    z = self.pool2(z)                        # [bs, 64, 4, 4]
   
    # neural network phase
    z = z.view(-1, 64 * 4 * 4)             # Size([bs, 1024])
    z = T.nn.functional.relu(self.fc1(z))  # Size([bs, 512])
    z = T.nn.functional.relu(self.fc2(z))  # Size([bs, 256])
    z = self.fc3(z)                        # Size([bs, 10])
    return z

def accuracy(model, ds):
  ldr = T.utils.data.DataLoader(ds,
    batch_size=len(ds), shuffle=False)
  n_correct = 0
  for data in ldr:
    (pixels, labels) = data
    with T.no_grad():
      oupts = model(pixels)
    (_, predicteds) = T.max(oupts, 1)
    n_correct += (predicteds == labels).sum().item()

  acc = (n_correct * 1.0) / len(ds)
  return acc

def main():
  # 0. setup
  print("\nBegin MNIST with CNN demo ")
  np.random.seed(1)
  T.manual_seed(1)

  # 1. create Dataset
  print("\nCreating 1000-item train Dataset from text file ")
  train_file = ".\\Data\\mnist_train_1000.txt"
  train_ds = MNIST_Dataset(train_file)

  bat_size = 10
  train_ldr = T.utils.data.DataLoader(train_ds,
    batch_size=bat_size, shuffle=True)

  # 2. create network
  print("\nCreating network with 2 conv and 3 linear ")
  net = Net().to(device)

  # 3. train model
  max_epochs = 50
  ep_log_interval = 5
  lrn_rate = 0.005

  loss_func = T.nn.CrossEntropyLoss()  # does log-softmax()
  optimizer = T.optim.SGD(net.parameters(), lr=lrn_rate)

  print("\nbat_size = %3d " % bat_size)
  print("loss = " + str(loss_func))
  print("optimizer = SGD")
  print("max_epochs = %3d " % max_epochs)
  print("lrn_rate = %0.3f " % lrn_rate)

  print("\nStarting training")
  net.train()  # set mode
  for epoch in range(0, max_epochs):
    epoch_loss = 0  # for one full epoch
    num_lines_read = 0  # allows quick exit during development

    for (batch_idx, batch) in enumerate(train_ldr):
      (X, y) = batch  # X = pixels, y = target labels
     
      num_lines_read += bat_size  # not exactly
      # if num_lines_read "gte" 10:
      #   continue

      optimizer.zero_grad()
      oupt = net(X)
      loss_val = loss_func(oupt, y)  # a tensor
      epoch_loss += loss_val.item()  # accumulate
      loss_val.backward()  # compute grads
      optimizer.step()     # update weights
    if epoch % ep_log_interval == 0:
      print("epoch = %4d   loss = %0.4f" % (epoch, epoch_loss))
  print("Done ") 

  # 4. evaluate model accuracy
  print("\nComputing model accuracy")
  net.eval()
  acc_train = accuracy(net, train_ds)  # all at once
  print("Accuracy on training data = %0.4f" % acc_train)

  test_file = ".\\Data\\mnist_test_100.txt"
  test_ds = MNIST_Dataset(test_file)
  net.eval()
  acc_test = accuracy(net, test_ds)  # all at once
  print("Accuracy on test data = %0.4f" % acc_test)

  # 5. make a prediction
  print("\nPredicting image for random 28x28 pixels ")
  digits = ['zero', 'one', 'two', 'three', 'four', 'five', 
    'six', 'seven', 'eight', 'nine' ]
  inpt = np.random.random((1, 1, 28, 28))  # 1 item, 1 channel
  inpt = T.tensor(inpt, dtype=T.float32).to(device)
  with T.no_grad():
    oupt = net(inpt)  # 10 logits like [[-0.12, 1.03, . . ]]
  am = T.argmax(oupt) # 0 to 9
  print("\nPredicted class is " + digits[am] )

  # 6. save model
  print("\nSaving trained model state")
  fn = ".\\Models\\mnist_model.pt"
  T.save(net.state_dict(), fn)  

  print("\nEnd MNIST CNN demo ")

if __name__ == "__main__":
  main()