MNIST Image Classification Using PyTorch 1.10 on Windows 11

One of my standard neural network examples is image classification on the MNIST dataset. The full MNIST (modified National Institute of Standards and Technology) dataset has 60,000 images for training and 10,000 images for testing. Each image is a 28 x 28 (784 pixels) grayscale handwritten digit from ‘0’ to ‘9’. Each pixel value is an integer from 0 (white) to 255 (black).

I fetched the raw MNIST data from http://yann.lecun.com/exdb/mnist/. The data is stored in four .gz (gnu-zipped) files: train-images-idx3-ubyte.gz, train-labels-idx1-ubyte.gz, t10k-images-idx3-ubyte.gz, t10k-labels-idx1-ubyte.gz. I used the 7-Zip utility program to extract the four files. The data is stored in a proprietary binary format so I wrote a helper program to convert the binary data to text files. See https://jamesmccaffreyblog.com/2022/02/25/preparing-mnist-image-data-text-files-in-visual-studio-magazine/.

I used a 1,000-item subset of the training data, and a 100-item subset of the test data. After conversion to text, the data looks like:

0 0 0 . . 84 185 159 . . 5
0 0 0 . . 133 254 87 . . 9
0 0 0 . . 164 79 202 . . 7
. . .

Each line is one image. The first 784 values on each line are the pixel values. The last value on each line is the target digit, ‘0’ to ‘9’.

I designed a convolutional neural network that has two convolution layers, three linear layers, two pooling layers, and two dropout layers. The architecture was adapted from an example I found buried in the PyTorch documentation.

I used ReLU() activation on all layers except for the final layer where I used no activation (combined with CrossEntropyLoss() for training which automatically adds log_softmax() activation). I used the default PyTorch weight and bias initialization. The documentation is not very clear about what this is, but I believe it’s xavier_uniform_() weight initialization and zeros_() bias initialization.

For training, I used stochastic gradient descent optimization with a fixed learning rate of 0.05 and a batch size of 20.

The demo achieved 97.00% accuracy on the test data: not bad but it’s possible to do better by fiddling with the hyperparameters.

To use the trained model, just for fun, I created a fake image that sort of resembles a mutated ‘4’ digit.



A mutated handwritten digit is one thing. A mutated plant in an old science fiction movie is another.

Left: “The Lost World” (1960) tells the story of explorers who find a hidden area on top of a nearly inaccessible plateau. There are several different types of plants that have mutated into carnivorous versions. The movie is loosely based on a novel by Arthur Conan Doyle, who best known for his Sherlock Holmes stories.

Center: “Matango” (1963) is a Japanese movie where survivors of a shipwreck end up on a deserted island. Unfortunately, eating mushrooms mutates people into mushroom-people.

Right: “The Day of the Triffids” (1962) tells a story where a global-wide meteor shower blinds most people on Earth and also delivers spores that mutate into huge Triffids — carnivorous plants that are slow moving but have a poisonous sting.


Demo code. Replace “lt”, “gt”, “lte”, “gte” with Boolean operator symbols. My lame blog editor chokes on symbols.

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

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

import numpy as np
import matplotlib.pyplot as plt
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):
    all_xy = np.loadtxt(src_file, usecols=range(785),
      delimiter="\t", comments="#", dtype=np.float32)

    tmp_x = all_xy[:, 0:784]  # all rows, cols [0,783]
    tmp_x /= 255.0
    tmp_x = tmp_x.reshape(-1, 1, 28, 28)  # bs, chnls, 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] 
    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.fc1 = T.nn.Linear(1024, 512)   # [64*4*4, x]
    self.fc2 = T.nn.Linear(512, 256)
    self.fc3 = T.nn.Linear(256, 10)     # 10 classes
    self.pool1 = T.nn.MaxPool2d(2, 2)   # kernel, stride
    self.pool2 = T.nn.MaxPool2d(2, 2)
    self.drop1 = T.nn.Dropout(0.25)
    self.drop2 = T.nn.Dropout(0.50)
    # uses default weight and bias initialization
  
  def forward(self, x):
    # convolution phase         # x is [bs, 1, 28, 28]
    z = T.relu(self.conv1(x))   # Size([bs, 32, 24, 24])
    z = self.pool1(z)           # Size([bs, 32, 12, 12])
    z = self.drop1(z)
    z = T.relu(self.conv2(z))   # Size([bs, 64, 8, 8])
    z = self.pool2(z)           # Size([bs, 64, 4, 4])
   
    # neural network phase
    z = z.reshape(-1, 1024)     # Size([bs, 1024])
    z = T.relu(self.fc1(z))     # Size([bs, 512])
    z = self.drop2(z)
    z = T.relu(self.fc2(z))     # Size([bs, 256])
    z = self.fc3(z)             # Size([bs, 10]) 
    return z  # implicit log-softmax() activation

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

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 PyTorch 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 = 20
  train_ldr = T.utils.data.DataLoader(train_ds,
    batch_size=bat_size, shuffle=True)

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

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

  # 3. train model
  max_epochs = 25  # 100 gives better results
  ep_log_interval = 5
  lrn_rate = 0.05
  
  loss_func = T.nn.CrossEntropyLoss()  # does log-softmax()
  optimizer = T.optim.SGD(net.parameters(), lr=lrn_rate)
  # optimizer = T.optim.Adam(net.parameters(), lr=0.005)
  
  print("\nbat_size = %3d " % bat_size)
  print("loss = " + str(loss_func))
  print("optimizer = SGD")
  print("lrn_rate = %0.3f " % lrn_rate)
  print("max_epochs = %3d " % max_epochs)


  print("\nStarting training")
  net.train()  # set mode
  for epoch in range(0, max_epochs):
    ep_loss = 0  # for one full epoch
    for (batch_idx, batch) in enumerate(train_ldr):
      (X, y) = batch  # X = pixels, y = target labels
      optimizer.zero_grad()
      oupt = net(X)
      loss_val = loss_func(oupt, y)  # a tensor
      ep_loss += loss_val.item()  # accumulate
      loss_val.backward()  # compute grads
      optimizer.step()     # update weights
    if epoch % ep_log_interval == 0:
      print("epoch = %4d   |  loss = %9.4f" % (epoch, ep_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. use model
  print("\nMaking prediction for fake image: ")
  x = np.zeros(shape=(28,28), dtype=np.float32)
  for row in range(5,23):
    x[row][9] = 180  # vertical line
  for rc in range(9,19):
    x[rc][rc] = 250  # diagonal
  for col in range(5,15):  
    x[14][col] = 200  # horizontal
  x /= 255.0

  plt.tight_layout()
  plt.imshow(x, cmap=plt.get_cmap('gray_r'))
  plt.show()

  x = x.reshape(1, 1, 28, 28)  # 1 image, 1 channel
  x = T.tensor(x, dtype=T.float32).to(device)
  with T.no_grad():
    oupt = net(x)  # 10 logits like [[-0.12, 1.03, . . ]]
  pred_probs = T.softmax(oupt, dim=1)
  print("\nPrediction probabilities: ")
  np.set_printoptions(formatter={'float': '{: 0.4f}'.format})
  # np.set_printoptions(precision=4, suppress=True)
  print(pred_probs.numpy())

  digits = ['zero', 'one', 'two', 'three', 'four', 'five', 
    'six', 'seven', 'eight', 'nine' ]
  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 PyTorch CNN demo ")

if __name__ == "__main__":
  main()
This entry was posted in Miscellaneous, PyTorch. Bookmark the permalink.