Session 21¶

Topic 3: Neural networks and PyTorch¶

• introduction to PyTorch
• automatic differentiation
• basic neural network models and training 👈
• optimization hyperparameters
• distributed training with data parallelism
• distributed training on multiple compute nodes

Today¶

Basic neural network models and training¶

1. building models with PyTorch
2. training loops in PyTorch

Building models in PyTorch¶

Tensors → gradients → models¶

so far:

• tensors … store data
• autograd … automatic differentiation

→ now: combine them into learnable models 🚀

What is a neural network model?¶

neural network model is a function with parameters which are learned from data

mathematically:

• \(x\) → input/features (from data)
• \(\theta\) → parameters (network weights & biases)
• \(\hat{y}\) → prediction (network output)

Example: 2-layer network¶

• \(W_1\in\mathbb{R}^{n \times d}, W_2\in\mathbb{R}^{d \times m}\) → weight matrices
• \(b_1\in\mathbb{R}^{d}, b_2\in\mathbb{R}^{m}\) → bias vectors
• \(\sigma\) → nonlinear function (activation) applied pointwise to a vector

What is theta exactly?¶

\(\theta\) refers to all weights and biases in the model:

Implementing a model in PyTorch¶

Key abstract class: nn.Module¶

all our models in PyTorch will inherit from

torch.nn.Module

the inheritance automatically gives us:

• parameter tracking → \(\theta\)
• autograd integration → \(\nabla_\theta\)
• composability with other models 🧩

The anatomy of a model¶

each model we write will have two parts:

1. __init__() → creation of layers
2. forward() → computation with the layers

Example:

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        # define all trainable layers (e.g., linear, convolution, ...)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # define computation f(x; θ) from input `x` to output `y`
        return y

Example: 2-layer dense neural network¶

import torch.nn as nn

class TwoLayerPerceptron(nn.Module):
    def __init__(self, in_features, out_targets):
        super().__init__()
        self.linear1 = nn.Linear(in_features, 4)
        self.activation = nn.ReLU()
        self.linear2 = nn.Linear(4, out_targets)

    def forward(self, x):
        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x

• The “2-layer” naming counts weight layers (\(W_1\) and \(W_2\)), not neuron layers — so there are actually three sets of nodes but only two sets of learnable parameters.
• ReLU is applied elementwise to the preactivation \(z = W_1 x + b_1\), zeroing any negative entries before they reach \(W_2\).
• The biases \(b_1\) and \(b_2\) are not drawn as explicit nodes here, but they are equivalent to adding a constant input of 1 with a learnable weight.

PyTorch tracks parameters automatically¶

model = TwoLayerPerceptron(3, 2)

for param in model.parameters():
    print(param)

• weights and biases are registered
• gradients are computed via autograd

→ no manual bookkeeping required 😌

Important distinction: parameter vs tensor¶

x = torch.zeros(1, 3)  # is a tensors

model = TwoLayerPerceptron(3, 2)
model.linear1.weight   # is a parameter

• a tensor is a generic multi-dimensional array
• a parameter is a trainable tensor

→ parameters automatically track gradients and get updated during training

Linear layers¶

(a.k.a. fully connected or dense)

linear = nn.Linear(in_features=3, out_features=4)

mathematically:

• \(W\) → weight matrix (here: \(4 \times 3\))
• \(b\) → bias vector (here: \(4\)-D)

How to use a linear layer?¶

x = torch.rand(1, 3)
z = linear(x)

• input: shape of x is (1, 3)
• output: shape of z is (1, 4)

Building deeper models¶

we concatenate layers:

input → linear layer → activation → linear layer → output

Activation functions – crucial¶

• without activation: model is a linear function only ❌
• with activation: model is a nonlinear function ✅

→ crucial for neural networks to approximate complex mappings (features → targets)

Common activation functions¶

In PyTorch:

torch.nn.ReLU()
torch.nn.LeakyReLU()
torch.nn.GELU()
torch.nn.SiLU()
torch.nn.Sigmoid()
torch.nn.Tanh()

Convolutional layers¶

often used for images:

conv = nn.Conv2d(3, 32, 5)

• input channels (e.g., \(3\) = RGB)
• output channels (e.g., \(32\))
• kernel size → weights are (just) a \(5 \times 5\) matrix

→ learns spatially local patterns

What is a convolution?¶

• sliding window over image
• detects patterns (edges, shapes)

→ like a filter — but learned on the data

Other layers¶

you should know more layers, but we skip them for now

• attention/transformer: learn correlations between features (what makes LLM possible)
• recurrent: for sequences (superseded by attention)
• residual: \(z = x + \sigma(W x + b)\)
• normalization: transform to zero mean, unit std.
• dropout: randomly set values to zero to prevent overfitting
• pooling: reduce image sizes (max/average over a small sliding window)

Putting it all together¶

a model is build with:

• trainable layers (Linear, Conv, Attention, …)
• activations (ReLU, GELU, SiLU …)
• normalization (LayerNorm GroupNorm, …)
• residuals ( x + layer(x) )
• regularization (Dropout)

Key takeaways¶

• nn.Module = foundation of all models
• \(\theta\) = all weights and biases
• define layers and \(\theta\) in   __init__()
• define \(f(x; \theta)\) in   forward()
• parameters are handled automatically

→ you’re now programming parameterized functions

Training loops in PyTorch¶

Model → learning¶

so far:

• tensors … store data
• autograd … automatic differentiation
• models … parameterized functions

→ now: make the model learn from data 🚀

What does “training” mean?¶

we want to find parameters \(\theta\) such that:

• predictions match data
• minimize the loss by changing \(\theta\):

→ training means solving an optimization problem

• we will use gradient descent:

The big picture¶

PyTorch training pipeline:

1. data → Dataset + DataLoader
2. model → neural network
3. loss → measure error (data vs model predictions)
4. optimizer → update parameters
5. loop → repeat until loss “looks good”

Data: Dataset & DataLoader¶

Why do we need them?¶

real datasets are:

• large 📦
• messy 🧹
• not in memory (to be loaded from disk)

Dataset¶

a Dataset gives structure

• returns one sample at a time

• defines:

__len__
__getitem__

💡 think: “index → (input, label)”

DataLoader¶

a DataLoader feeds the training loop

• wraps a Dataset

• provides:

• batching 📚

• shuffling 🔀
• parallel loading ⚡

Example: MNIST¶

import torch
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

training_set = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

training_loader = torch.utils.data.DataLoader(
    training_set, batch_size=256, shuffle=True
)

Loss Function¶

What is loss?¶

a scalar-valued function

that measures how wrong the model is

Example: cross-entropy loss¶

loss_fn = torch.nn.CrossEntropyLoss()

y_pred = model(x)  # logits: raw, unnormalized scores from the model
loss_fn(y_true, y_pred)

used for:

• classification problems
• logits (not probabilities)

logits:

• can be a vector with any real numbers (e.g. [-1.2, 0.3, 2.1])
• do not sum to 1
• softmax is applied by the cross-entropy loss

Cross-entropy intuition¶

• correct class → high score → low loss ✅
• wrong class → high score → high loss ❌

💡 pushes probability mass toward the correct class

When to use cross-entropy?¶

• image classification (e.g. MNIST, CIFAR-10) 🖼️
• text classification (spam detection, sentiment) 💬
• multi-class prediction problems

Example: mean squared error (MSE)¶

loss_fn = torch.nn.MSELoss()

y_pred = model(x)  # raw output vectors
loss_fn(y_pred, y_true)

used for:

• regression problems
• predicting continuous values

MSE intuition¶

• penalizes squared differences
• large errors → punished heavily

💡 encourages predictions to be close in magnitude

When to use MSE?¶

• predicting prices 💰
• forecasting temperatures 🌡️
• any continuous output

Optimizer¶

What does the optimizer do?¶

updates parameters using gradients:

\(\eta > 0\) is called learning rate

Example: SGD with momentum¶

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.001,
    momentum=0.9
)

hyperparameters:

• learning rate ( lr ) → step size
• momentum → smooth updates across data samples

⚠️ needs tuning: wrong learning rate/momentum → no learning

Example: Adam optimizer¶

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999)
)

hyperparameters:

• learning rate ( lr ) → step size
• betas → control moving averages of gradients

🧠 adapts the learning rate per parameter

⚠️ often works well with defaults, but still needs tuning

Why Adam is a good default choice?¶

• combines momentum and adaptive scaling
• handles noisy / sparse gradients well
• typically converges faster than SGD

The training loop¶

Core idea¶

sequence inside the training loop:

1. get batch
2. clear old gradients (recall autograd behavior)
3. compute prediction (forward pass)
4. compute loss
5. compute gradients (backward pass)
6. update parameters

One training step in PyTorch¶

inputs, labels = data             # 1. get batch

optimizer.zero_grad()             # 2. clear old gradients
outputs = model(inputs)           # 3. compute prediction
loss = loss_fn(outputs, labels)   # 4. compute loss
loss.backward()                   # 5. compute gradients
optimizer.step()                  # 6. update parameters

→ this is gradient descent in code

Training over batches¶

for i, data in enumerate(training_loader):
    # perform one training step

• loop over dataset in batches
• each batch = one update

→ this is stochastic gradient descent in code because DataLoader is shuffeling the order of samples

Training over epochs¶

for epoch in range(num_epochs):
    for i, data in enumerate(training_loader):
        # perform one training step

• one epoch: one full pass over the dataset
• DataLoader is shuffling the samples from one epoch to another

Validation 📊¶

Why validate?¶

we actually care about:

• generalization (samples the model did not train on)
• not memorization (good predictions for training samples only)

Validation loop:

• run model on unseen data, log performance
• do not update weights with validation samples

The full training pipeline¶

→ this loop at is the heart of deep learning 🎯

Common pitfalls¶

• forgetting zero_grad()
• wrong loss function
• learning rate too high/low
• not using validation
• overfitting because of too many epochs

not preprocessing the data ⚠️

• inconsistent scales → features with large magnitudes dominate learning (unstable gradients)
• raw / unnormalized data → slower convergence or no learning (model struggles to find useful patterns)

Key takeaways¶

• training = optimize parameters via gradients
• DataLoader → feeds batches
• loss → defines objective
• optimizer → updates weights
• training loop → ties everything together

🎯 If you understand this loop, you understand PyTorch.

Looking ahead¶

now we ask: why does training sometimes work…
and sometimes fail?

→ optimization hyperparameters

Prepare for the next session 🌱¶

sessions/22-todo