PyTorch — Module 3.3

How This Lesson Fits the Module

Scikit-learn handles classical ML with clean pipelines. PyTorch is where most modern AI engineering happens: neural networks, transformers, diffusion models, and custom training loops with automatic differentiation.

PyTorch’s imperative style—build the graph as you run code—matches how engineers debug. Research labs, startups, and an increasing share of production systems standardize on PyTorch. Compare with TensorFlow to choose the right stack for your team.

Learning Objectives

By the end of this lesson, students should be able to:

Explain tensors, autograd, and the training loop in PyTorch.
Build a nn.Module and run forward/backward passes.
Load data with DataLoader and train on CPU or GPU.
Save and load model checkpoints for reproducibility.
Decide when PyTorch is the right framework for a project.
Recognize how PyTorch integrates with the Hugging Face ecosystem.

What PyTorch Is—and When to Use It

PyTorch is an open-source deep-learning framework centered on GPU-accelerated tensors and dynamic computation graphs. Its Pythonic API makes experimentation fast; torch.compile and TorchScript bridge to production performance.

Choose PyTorch when…	Consider TensorFlow when…
Research, prototyping, and custom architectures dominate	Your org standardizes on TFX / Google Cloud ML
You need Hugging Face Transformers natively	You deploy extensively with TensorFlow Lite on mobile/edge
Debugging with standard Python tools is a priority	Legacy production graphs are already in TensorFlow
Dynamic control flow (variable-length sequences) is common	You want Keras’s high-level API as the primary interface

Tensors and Autograd

PyTorch tensors generalize NumPy arrays with GPU support and gradient tracking. requires_grad=True tells autograd to record operations for backpropagation.

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(32, 784, device=device, requires_grad=True)
w = torch.randn(784, 128, device=device, requires_grad=True)
b = torch.zeros(128, device=device, requires_grad=True)

h = torch.relu(x @ w + b)
loss = h.mean()
loss.backward()   # computes gradients for w, b
print(w.grad.shape)  # (784, 128)

Definition — Training Loop

A training loop repeats: forward pass → loss computation → backward() → optimizer step → zero_grad(). Every deep-learning framework implements this cycle; PyTorch exposes it explicitly so you control every detail.

Defining a Model with nn.Module

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim=784, hidden=128, n_classes=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden, n_classes),
        )

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

model = MLP().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

DataLoader and a Complete Training Step

from torch.utils.data import DataLoader, TensorDataset

dataset = TensorDataset(X_tensor, y_tensor)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

model.train()
for batch_x, batch_y in loader:
    batch_x, batch_y = batch_x.to(device), batch_y.to(device)

    logits = model(batch_x)
    loss = criterion(logits, batch_y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

ML Example — Transfer Learning

Load a pretrained ResNet from torchvision.models, freeze backbone layers, and fine-tune the classifier head on your dataset. This pattern—pretrained backbone + small custom head—powers most computer-vision production systems.

Inference and Checkpointing

# Save
torch.save({
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "epoch": epoch,
}, "checkpoint.pt")

# Load
ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])

model.eval()
with torch.no_grad():
    preds = model(batch_x).argmax(dim=1)

PyTorch Strengths

Pythonic, debuggable imperative API
Dominant in research and LLM tooling
Strong GPU/TPU support via ecosystem packages
Composable with NumPy, sklearn preprocessing

PyTorch Trade-offs

Explicit loops require discipline (AMP, grad clipping)
Deployment needs extra tooling (TorchServe, ONNX)
Large models demand significant GPU memory
API surface grows fast—pin versions in production

Common Misconception: “PyTorch replaces sklearn for tabular data.”

Reality: Gradient-boosted trees and logistic regression often beat small neural nets on structured data. Use PyTorch when representation learning, scale, or unstructured modalities justify the complexity.

Knowledge Check

Short Answer: What does loss.backward() compute? Answer: Gradients of the loss with respect to parameters that require grad.
True/False: model.eval() and torch.no_grad() are needed for inference. Answer: True (best practice).
Short Answer: Why call optimizer.zero_grad() each step? Answer: PyTorch accumulates gradients by default; zero them before the next backward pass.
Multiple Choice: Primary Hugging Face backend: (a) TensorFlow, (b) PyTorch, (c) JAX only. Answer: (b).
Short Answer: What does requires_grad=True tell autograd? Answer: Record operations so gradients can be computed during backward().
True/False: Custom models typically subclass nn.Module and implement forward. Answer: True.
Short Answer: Why call DataLoader(..., shuffle=True) during training? Answer: Shuffle batches each epoch to reduce order bias and improve SGD noise.
Multiple Choice: Checkpoints usually save: (a) only print statements, (b) state_dict of model (and often optimizer), (c) the entire OS, (d) HTML. Answer: (b).
True/False: PyTorch should replace sklearn for every tabular problem. Answer: False—trees and linear models often win on structured data.
Short Answer: List the training-loop steps in order. Answer: Forward pass, compute loss, optimizer.zero_grad(), loss.backward(), optimizer.step().

Key Takeaways

PyTorch is the leading framework for neural network research and LLM engineering.
Tensors + autograd + explicit training loops give full control.
Use nn.Module, DataLoader, and checkpoints as core patterns.
Match the tool to the problem—not every task needs a neural network.
Next: TensorFlow for an alternative production-oriented stack.

Trainer’s Guide

Hands-on idea: MNIST classifier in PyTorch with logged train/val loss curves plotted in Matplotlib. Require GPU detection, checkpoint saving, and a final test accuracy report.

Recap: PyTorch combines tensors, autograd, and explicit training loops; next, compare the Keras/deployment-oriented stack in TensorFlow.