PyTorch Basics

In the last lesson you calculated gradients by hand. Deep learning frameworks do that for you and run the calculations on graphics processors. PyTorch is the most widely used framework in research and is common in industry. TensorFlow with Keras and JAX are the main alternatives.

PyTorch does not run in the browser, so use your own computer or a free hosted notebook such as Google Colab. Install it with pip install torch.

Tensors

A tensor is PyTorch’s array type. It works like a NumPy array, with two additions: it can live on a GPU, and it can track gradients.

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x.shape, x.dtype)
print(x * 2)
print(x @ x)              # matrix multiplication
print(x.mean(), x.sum(dim=0))

device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)

Automatic differentiation

Mark a tensor with requires_grad=True and PyTorch records every operation on it. Calling backward() on a result calculates the gradients.

import torch

w = torch.tensor(-4.0, requires_grad=True)
loss = (w - 3) ** 2 + 2
loss.backward()

print(w.grad)   # tensor(-14.), the same as 2 * (w - 3)

This is the gradient you wrote by hand in the previous lesson.

Defining a model

A model is a class that inherits from nn.Module. You create the layers in __init__ and describe the forward pass in forward.

import torch
from torch import nn

class Classifier(nn.Module):
    def __init__(self, n_features, n_classes):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(n_features, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, n_classes),
        )

    def forward(self, x):
        return self.layers(x)

model = Classifier(n_features=64, n_classes=10)
print(model)
print(sum(p.numel() for p in model.parameters()), "parameters")

nn.Linear is a fully connected layer. The model returns raw scores, called logits. The loss function applies softmax internally.

Data loading

DataLoader serves the data in shuffled batches.

import torch
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

X, y = load_digits(return_X_y=True)
X_train, X_val, y_train, y_val = train_test_split(X / 16.0, y, random_state=0, stratify=y)

to_tensor = lambda a, kind: torch.tensor(a, dtype=kind)
train_data = TensorDataset(to_tensor(X_train, torch.float32), to_tensor(y_train, torch.long))
val_data = TensorDataset(to_tensor(X_val, torch.float32), to_tensor(y_val, torch.long))

train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
val_loader = DataLoader(val_data, batch_size=256)

The training loop

Every PyTorch project contains a version of this loop. It continues from the two blocks above.

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(20):
    model.train()
    for features, labels in train_loader:
        optimizer.zero_grad()            # clear old gradients
        loss = loss_fn(model(features), labels)
        loss.backward()                  # calculate gradients
        optimizer.step()                 # update the weights

    model.eval()
    correct = 0
    with torch.no_grad():                # no gradients needed for evaluation
        for features, labels in val_loader:
            correct += (model(features).argmax(dim=1) == labels).sum().item()
    print(f"epoch {epoch + 1}: validation accuracy {correct / len(val_data):.3f}")

The five steps inside the inner loop are the ones to remember:

1. zero_grad() clears the gradients from the previous batch.
2. The forward pass produces predictions.
3. The loss function compares them with the labels.
4. backward() calculates the gradients.
5. step() updates the weights.

train and eval modes

model.train() turns on behaviour used only in training, such as dropout. model.eval() turns it off. Forgetting to switch is a common source of confusing results.

Saving and loading

torch.save(model.state_dict(), "digits.pt")

model = Classifier(n_features=64, n_classes=10)
model.load_state_dict(torch.load("digits.pt"))
model.eval()

Practical advice

• Start with a small model and a small sample of data, and confirm that the loss falls.
• Check that the model can overfit a handful of examples. If it cannot, there is a bug.
• Plot training and validation loss together.
• Set random seeds so that runs can be compared.
• Use a GPU for anything larger than toy data.

Practice

Run the training loop, then remove the nn.Dropout(0.2) layer and run it again. Compare the validation accuracy.