PyTorch Tutorial

A working mental model of how deep learning libraries actually work, followed by a complete training loop you can copy and adapt.

1. The mental model

A neural network is just a function that maps an input tensor to an output tensor. Training is the process of finding the parameters of that function so that, on average, the output is close to the truth.

Every deep learning library — PyTorch, TensorFlow, JAX — has four pieces:

1. Tensors — n-dimensional arrays that live on a CPU or GPU.
2. Autograd — bookkeeping that lets the library compute the gradient of any scalar with respect to any tensor that touched it.
3. Modules (models) — reusable bundles of tensors and operations (convolutions, ReLUs, attention, ...). A model is itself a Module.
4. Optimizers — rules for updating tensors using their gradients (SGD, Adam, AdamW).

2. Tensors

import torch

# Create a tensor from a Python list
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
print(x.shape)        # torch.Size([2, 2])
print(x.dtype)        # torch.float32

# Move it to GPU if you have one
device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()  # Apple Silicon
    else "cpu"
)
x = x.to(device)

# An image batch has shape (N, C, H, W):
#   N = batch size, C = channels (3 for RGB), H/W = pixels
batch = torch.randn(8, 3, 224, 224, device=device)

3. Autograd in one example

w = torch.tensor(2.0, requires_grad=True)   # a parameter we want to optimize
x = torch.tensor(3.0)
y_true = torch.tensor(10.0)

y_pred = w * x          # forward pass
loss = (y_pred - y_true) ** 2   # squared error
loss.backward()         # PyTorch fills in w.grad
print(w.grad)           # d(loss)/d(w) = 2 * (w*x - y_true) * x = -24

# A single SGD step: nudge w in the direction that reduces loss.
with torch.no_grad():
    w -= 0.01 * w.grad
    w.grad.zero_()

That's the entire idea. A real training loop just does this billions of times with a much bigger w (millions of weights) and many examples in a batch.

4. Datasets and DataLoaders

A Dataset tells PyTorch how to get one example. A DataLoader wraps a Dataset to give you batches, in parallel, with shuffling.

from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms
from pathlib import Path

class CameraTrapDataset(Dataset):
    """One example = one image + its label.

    `root/day/*.JPG` are day images (label 0); `root/night/*.JPG` are night (label 1).
    """
    def __init__(self, root, transform):
        self.items = []
        for label, sub in enumerate(["day", "night"]):
            for p in (Path(root) / sub).glob("*.JPG"):
                self.items.append((p, label))
        self.transform = transform

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

    def __getitem__(self, idx):
        path, label = self.items[idx]
        img = Image.open(path).convert("RGB")
        img = self.transform(img)     # convert to tensor, resize, normalize
        return img, label

# Standard ImageNet preprocessing
tx = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

train_ds = CameraTrapDataset("data/train", transform=tx)
val_ds   = CameraTrapDataset("data/val",   transform=tx)

train_dl = DataLoader(train_ds, batch_size=32, shuffle=True,  num_workers=4)
val_dl   = DataLoader(val_ds,   batch_size=32, shuffle=False, num_workers=4)

5. A model

For everything in this course you'll start from a model that someone else trained on ImageNet, then fine-tune the last layer. That gets you 90% of the way for a few percent of the compute.

import torch.nn as nn
from torchvision import models

# Load a ResNet-18 with ImageNet weights
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)

# Replace the final 1000-class classifier with a 2-class one (day vs night)
model.fc = nn.Linear(model.fc.in_features, 2)
model = model.to(device)

6. The training loop

Every PyTorch training script in existence looks like this. Learn this loop and you can read almost any training code.

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4)

def run_one_epoch(loader, train):
    model.train(train)
    total, correct, loss_sum = 0, 0, 0.0
    for imgs, labels in loader:
        imgs   = imgs.to(device)
        labels = labels.to(device)

        with torch.set_grad_enabled(train):
            logits = model(imgs)                # forward
            loss   = criterion(logits, labels)  # how wrong were we?

            if train:
                optimizer.zero_grad()           # clear old gradients
                loss.backward()                 # compute new gradients
                optimizer.step()                # update the weights

        preds = logits.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total   += labels.size(0)
        loss_sum += loss.item() * labels.size(0)
    return loss_sum / total, correct / total

n_epochs = 5
for epoch in range(n_epochs):
    train_loss, train_acc = run_one_epoch(train_dl, train=True)
    val_loss,   val_acc   = run_one_epoch(val_dl,   train=False)
    print(f"epoch {epoch+1}: "
          f"train_loss={train_loss:.3f} train_acc={train_acc:.3f}  "
          f"val_loss={val_loss:.3f}  val_acc={val_acc:.3f}")

7. Saving and loading

# Save only the weights (the recommended pattern)
torch.save(model.state_dict(), "day_vs_night_resnet18.pt")

# Loading later, in a fresh script
model = models.resnet18()
model.fc = nn.Linear(model.fc.in_features, 2)
model.load_state_dict(torch.load("day_vs_night_resnet18.pt", map_location=device))
model.eval()

8. Inference on a single image

img = Image.open("IMG_0042.JPG").convert("RGB")
x = tx(img).unsqueeze(0).to(device)   # add the batch dimension

with torch.no_grad():
    logits = model(x)
    probs  = logits.softmax(dim=1)[0].cpu().numpy()

print({"day": probs[0], "night": probs[1]})

9. Things that will trip you up

→ Continue to Lab 2: Train day-vs-night on Oregon Critters