PyTorch Basics

I'll assume that everyone is familiar with Python. Training neural nets in bare Python is somewhat painful, but fortunately there are several well-established libraries which can help. I like PyTorch, which is built upon an earlier library called Torch. (There are many others, including TensorFlow and Jax, both of which I would encourage you to explore at your leisure.)

import numpy as np
import torch
import torchvision

Torch believes that everything is a tensor.

a = np.random.rand(2,3)
b = torch.from_numpy(a)

print(a)
print(b)

[[0.5656557  0.27607714 0.45466866]
 [0.50539309 0.87329605 0.83471437]]
tensor([[0.5657, 0.2761, 0.4547],
        [0.5054, 0.8733, 0.8347]], dtype=torch.float64)

The main intuition is that tensors allow for intuitive and efficient matrix multiplication across different indexing dimensions. Soon we will see that training neural nets basically involves forward and backward passes, both of which are essentially matrix multiplies.

The other thing about Torch variables is that they (natively) can be differentiated. (Again, we will see why this is important when we learn about backpropagation in deep learning.)

A = torch.rand(2,2)
b = torch.rand(2,1)
x = torch.rand(2,1,requires_grad=True)

y = torch.matmul(A,x) + b

print(y)

tensor([[1.1021],
        [1.4504]], grad_fn=<AddBackward0>)

z = y.sum()
print(z)

tensor(2.5525, grad_fn=<SumBackward0>)

Here, z is a function of the input x. Let us now compute the derivative of z with respect to x using backpropagation.

z.backward()
print(x.grad)
print(x)

tensor([[1.2117],
        [1.4395]])
tensor([[0.1127],
        [0.7029]], requires_grad=True)

Training simple models

OK, enough background. Let us now train a simple model. We will train a logistic classifier (equivalent to a single-layer neural network) on a popular image dataset called Fashion-MNIST. (TorchVision has several other image datasets which we can directly load as variables, feel free to train your own models on them).

trainingdata = torchvision.datasets.FashionMNIST('./FashionMNIST/',train=True,download=True,transform=torchvision.transforms.ToTensor())
testdata = torchvision.datasets.FashionMNIST('./FashionMNIST/',train=False,download=True,transform=torchvision.transforms.ToTensor())

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to ./FashionMNIST/FashionMNIST/raw/train-images-idx3-ubyte.gz

100%|██████████████████████████| 26421880/26421880 [00:04<00:00, 5676746.07it/s]

Extracting ./FashionMNIST/FashionMNIST/raw/train-images-idx3-ubyte.gz to ./FashionMNIST/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to ./FashionMNIST/FashionMNIST/raw/train-labels-idx1-ubyte.gz

100%|████████████████████████████████| 29515/29515 [00:00<00:00, 1842706.76it/s]

Extracting ./FashionMNIST/FashionMNIST/raw/train-labels-idx1-ubyte.gz to ./FashionMNIST/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to ./FashionMNIST/FashionMNIST/raw/t10k-images-idx3-ubyte.gz

100%|████████████████████████████| 4422102/4422102 [00:00<00:00, 5298891.24it/s]

Extracting ./FashionMNIST/FashionMNIST/raw/t10k-images-idx3-ubyte.gz to ./FashionMNIST/FashionMNIST/raw

Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz
Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to ./FashionMNIST/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz

100%|█████████████████████████████████| 5148/5148 [00:00<00:00, 24206588.56it/s]

Extracting ./FashionMNIST/FashionMNIST/raw/t10k-labels-idx1-ubyte.gz to ./FashionMNIST/FashionMNIST/raw

Let's quickly check that everything has been downloaded.

print(len(trainingdata))
print(len(testdata))

60000
10000

Some of you might know that this is the same size as the (very commonly used) MNIST dataset. Let's plot some images.

image, label = trainingdata[0]
print(image.shape, label)

torch.Size([1, 28, 28]) 9

We cannot directly plot the 'image' object since it is a Torch tensor, so let's convert it back into a numpy array before displaying it. We will use matplotlib to show images.

import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(image.squeeze().numpy())
plt.show()

OK, that kinda sorta looks like a shoe? FashionMNIST is basically a bunch of black-white thumbnail images of different pieces of clothing. Let's try and plot a few of them. To step through the dataset, it will be helpful to construct a data loader object; we will use this while training our model below as well.

trainDataLoader = torch.utils.data.DataLoader(trainingdata,batch_size=64,shuffle=True)
testDataLoader = torch.utils.data.DataLoader(testdata,batch_size=64,shuffle=False)

Let's now use the DataLoader object to spit out a few images from the dataset.

images, labels = next(iter(trainDataLoader))

plt.figure(figsize=(10,4))
for index in np.arange(5):
    plt.subplot(1, 5, index + 1)
    # Assuming the images are in the format [C, H, W] and are grayscale
    # Adjust the image for visualization (transpose axes from [C, H, W] to [H, W, C])
    # For grayscale images, we also need to remove the channel dimension using squeeze()
    image = images[index].squeeze().numpy() # Assuming the images are tensors
    plt.imshow(image, cmap='gray')
    plt.axis('off')  # Optionally remove axes
plt.show()

OK! Time to set up our model.

class LinearReg(torch.nn.Module):
  def __init__(self):
    super(LinearReg, self).__init__()
    self.linear = torch.nn.Linear(28*28,10)

  def forward(self, x):
    x = x.view(-1,28*28)
    transformed_x = self.linear(x)
    return transformed_x

#net = LinearReg().cuda()
net = LinearReg()

Loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)

Cool! Everything is set up. Let's now train our network.

train_loss_history = []
test_loss_history = []

for epoch in range(20):
  train_loss = 0.0
  test_loss = 0.0
  for i, data in enumerate(trainDataLoader):
    images, labels = data
    images = images #.cuda()
    labels = labels #.cuda()
    optimizer.zero_grad()
    predicted_output = net(images)
    fit = Loss(predicted_output,labels)
    fit.backward()
    optimizer.step()
    train_loss += fit.item()
  for i, data in enumerate(testDataLoader):
    with torch.no_grad():
      images, labels = data
      images = images#.cuda()
      labels = labels#.cuda()
      predicted_output = net(images)
      fit = Loss(predicted_output,labels)
      test_loss += fit.item()
  train_loss = train_loss/len(trainDataLoader)
  test_loss = test_loss/len(testDataLoader)
  train_loss_history.append(train_loss)
  test_loss_history.append(test_loss)
  print('Epoch %s, Train loss %s, Test loss %s'%(epoch, train_loss, test_loss))

Epoch 0, Train loss 0.6664804585897592, Test loss 0.6461255649092851
Epoch 1, Train loss 0.6027035086965764, Test loss 0.6020751965653365
Epoch 2, Train loss 0.5679286457836501, Test loss 0.5770273548402604
Epoch 3, Train loss 0.5455557886979727, Test loss 0.5585913810001057
Epoch 4, Train loss 0.5292342450064637, Test loss 0.545891882128017
Epoch 5, Train loss 0.5165409704706053, Test loss 0.5348615168006556
Epoch 6, Train loss 0.506088534477296, Test loss 0.5280016594252009
Epoch 7, Train loss 0.49766681459285556, Test loss 0.5239671259928661
Epoch 8, Train loss 0.49072493615943485, Test loss 0.5145675993649064
Epoch 9, Train loss 0.4845343728277729, Test loss 0.5095292973290583
Epoch 10, Train loss 0.4793138336588833, Test loss 0.5043158817822766
Epoch 11, Train loss 0.4742559713087102, Test loss 0.5014074294810082
Epoch 12, Train loss 0.4701572920976163, Test loss 0.4970778027537522
Epoch 13, Train loss 0.46619399743420736, Test loss 0.4978971555354489
Epoch 14, Train loss 0.46299884320576307, Test loss 0.49122682213783264
Epoch 15, Train loss 0.45990598938866717, Test loss 0.4904827428091863
Epoch 16, Train loss 0.45701435573700905, Test loss 0.4862680175122182
Epoch 17, Train loss 0.45433947061107105, Test loss 0.4846375390982172
Epoch 18, Train loss 0.45184447986484844, Test loss 0.484140664718713
Epoch 19, Train loss 0.4495418051730341, Test loss 0.4815634836436837

Let's see how we did! We have tracked the losses so let's plot it.

plt.plot(range(20),train_loss_history,'-',linewidth=3,label='Train error')
plt.plot(range(20),test_loss_history,'-',linewidth=3,label='Test error')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.grid(True)
plt.legend()

<matplotlib.legend.Legend at 0x7f16dd79f130>

Not too bad (although more epochs would have helped, I suppose). Let's evaluate our final model on a few more images to wrap things up.

predicted_output = net(images)
print(torch.max(predicted_output, 1))
fit = Loss(predicted_output, labels)
print(labels)
print(fit)

torch.return_types.max(
values=tensor([ 6.2214,  3.0903,  9.0717,  7.8366,  6.7475,  6.2147, 10.1018,  4.5116,
         7.1518, 11.8602, 10.5575, 10.6477,  6.8496,  4.6596,  9.4688,  4.5188],
       grad_fn=<MaxBackward0>),
indices=tensor([3, 1, 7, 5, 8, 2, 5, 6, 8, 9, 1, 9, 1, 8, 1, 5]))
tensor([3, 2, 7, 5, 8, 4, 5, 6, 8, 9, 1, 9, 1, 8, 1, 5])
tensor(0.3164, grad_fn=<NllLossBackward0>)

Cool, looks like 2 images in this batch are being misclassified; can you (as a human) guess what those are?

plt.figure(figsize=(10,10))
for index in np.arange(0,16):
  plt.subplot(4,4,index+1)
  plt.imshow(images[index].cpu().squeeze().numpy(),cmap=plt.cm.gray)