Implementing Neural Network Training in PyTorch

Last updated October 16, 2026

David I. Inouye

Implementation connects the graph to running code

Arc 5 · Implementation in PyTorch

1. Autograd: inspect the scalar computation and its gradients.
2. Training: assemble the image model, loss, and optimizer.
3. Diagnostics: inspect what the training loop produces.

The preceding section explained backpropagation mathematically. Now we make each operation concrete.

30-second orientation. This is a distinct implementation section after mathematical backpropagation, not a new derivation. Walk through one persistent scalar example, then reuse the same pattern for the Fashion-MNIST network. All code is ordinary eager-mode CPU PyTorch; no compiler or GPU prerequisite. Labeled outputs were captured by executing the displayed snippets in slide order with PyTorch 2.7.1; the companion script reproduces that order. Numerical training outputs can vary by environment.

requires_grad identifies which inputs need derivatives

For scalar \(b=f(\boldsymbol a)\), computing \(\nabla_{\boldsymbol a} b\) requires marking a with requires_grad=True.

import torch
x = torch.tensor(2.)
y = torch.tensor(1.)
w = torch.tensor(1., requires_grad=True)
v = torch.tensor(1., requires_grad=True)
print(x.requires_grad, w.requires_grad, w.grad)

Output:

False True None

Mark the variable in the subscript of \(\nabla\)—or the denominator of \(\partial b/\partial a\) for scalars.

Here w and v are the trainable leaves. Their .grad fields will hold the weight derivatives after backward.

About three minutes. Connect the differentiation variable in the notation to requires_grad on that tensor, before performing the forward operations. For a scalar output and vector input, the gradient has one entry per input feature; for scalar input a, the notation is partial b / partial a. The scalar example below requests weight derivatives, so fixed data x and y need not require gradients. Introduce a tensor as a numerical array; these are zero-dimensional floating-point tensors. requires_grad marks inputs for differentiation in normal grad mode. None means no gradient has been stored yet, not that the derivative is zero. Parameters can receive gradients even when data do not require them. An operation can still save an untracked input value needed for a tracked weight’s derivative. [Sources] source:pytorch-autograd-2-7.

The forward pass builds values and a backward graph

print("Before computation, w.grad_fn:", w.grad_fn)
z = w * x
print("After computation, z.grad_fn:", type(z.grad_fn).__name__)
h = torch.relu(z)
s = v * h
r = s - y
loss = 0.5 * r.square()
print(z.item(), h.item(), s.item(), loss.item())

Output:

Before computation, w.grad_fn: None
After computation, z.grad_fn: MulBackward0
2.0 2.0 2.0 0.5

Forward values: \(w,x\to z\to h\); then \(v,h\to s\to r\to\ell\).

Backward graph: w is a leaf with no producing operation. Multiplying it creates z with a .grad_fn linking to that operation.

requires_grad=True enables tracking; the forward computation creates the graph.

About three minutes. Map each line to the preceding theoretical example. The graph follows operations actually executed, rather than parsing a symbolic expression in advance. The result carries a grad_fn that links into the backward graph. The loss here is a scalar, so backward will start with a seed of one. .item is used only for displaying scalar values. [Sources] source:pytorch-autograd-2-7.

Backward rules need selected forward values

For our scalar example:

Operation	A value its backward rule needs	Here
\(z=wx\)	\(x\) to obtain \(\partial\ell/\partial w\)	\(2\)
\(h=\operatorname{ReLU}(z)\)	Which inputs were positive	Active
\(s=vh\)	\(h\) for \(v\)’s derivative; \(v\) for \(h\)’s	\(2\), \(1\)
\(\ell=\tfrac12 r^2\)	Residual \(r\)	\(1\)

Saved forward values are not gradients. They let backward evaluate the local derivative rules.

About three minutes. Connect the table directly to the preceding mathematical worked example, without re-deriving the full chain. This is a list of mathematically sufficient information, not a promise that every implementation stores this exact minimal representation. PyTorch may save a ReLU output rather than a Boolean mask. The next slide observes the actual saved tensors. [Sources] source:pytorch-autograd-2-7.

We can observe what autograd saves for backward

saved = []
def pack(t):
    saved.append(t.detach().clone())  # snapshot for inspection
    return t.detach()                # value kept for backward

def unpack(t):
    return t

with torch.autograd.graph.saved_tensors_hooks(pack, unpack):
    loss = 0.5 * (v * torch.relu(w * x) - y).square()
print([t.item() for t in saved])

Output:

[2.0, 2.0, 2.0, 1.0, 1.0]

These values supply the multiplication, ReLU, and squared-loss backward rules.

About four minutes including reading the output. This is an observation tool, not code students need in every training loop. In the tested build the sequence is x, the ReLU output, h, v, and the residual. Repeated values can serve different operations; this log does not establish distinct storage allocations. pack receives a tensor being saved; unpack returns it when backward needs it. detach removes recorded history, not the numerical value; cloning the observation makes a stable snapshot. Exact saving order and representation may change by version. The hooks do not modify values and backward must still produce the same gradients. [Sources] source:pytorch-autograd-2-7.

One backward() computes both weight gradients

print("Before backward:", w.grad, v.grad)
loss.backward()
print("After backward:", w.grad.item(), v.grad.item())
print("Weights:", w.item(), v.item())

Output:

Before backward: None None
After backward: 2.0 2.0
Weights: 1.0 1.0

Backward starts at the scalar loss, follows the recorded graph, and adds contributions into the trainable leaves’ .grad fields.

The weights have not changed. Ordinary backward releases saved values; recompute forward for the next batch or updated weights.

About three minutes. Verify against the earlier hand calculation: both parameter derivatives are two. The saved-tensor hook example computed the same loss with the same numerical values. Backward uses those saved values and propagates sensitivities in reverse. It does not use finite perturbations. For a non-scalar output one must specify the desired upstream vector or reduce to a scalar; training here uses a scalar mean loss. [Sources] source:pytorch-autograd-2-7.

What does .grad contain on the second iteration?

a = torch.tensor(1., requires_grad=True)
for iteration in range(2):
    b = a.square()
    b.backward()
    print(iteration + 1, a.grad.item())

Output after iteration 1:

1 2.0

What value of a.grad.item() will iteration 2 print?

A. \(2.0\)   B. \(4.0\)   C. \(0.0\)

Output after iteration 2:

2 4.0

Each backward() adds to the previous .grad: \(2+2=4\). To get fresh gradients, reset them before recomputing a gradient.

About three minutes including voting. Do not explain accumulation or clearing before the vote. This self-contained example creates a fresh graph each iteration and never changes a. Each derivative is two, but the stored total becomes four. Reveal the second output and explain that backward adds to existing gradients; the next slide introduces how to reset them. [Sources] source:pytorch-autograd-2-7.

zero_grad() begins a new gradient accumulation

def objective():
    return 0.5 * (v * torch.relu(w * x) - y).square()

optimizer = torch.optim.SGD([w, v], lr=0.1)
optimizer.zero_grad()
print("After reset:", w.grad)
loss = objective()
loss.backward()
print("After backward:", w.grad.item())

Output:

After reset: None
After backward: 2.0

zero_grad() clears the previous accumulated gradients.

It resets gradients, not parameter values. Clear before each independent batch’s backward pass.

About three minutes. SGD is the already-taught gradient descent update; choose it for transparent scalar arithmetic. The optimizer holds references to the supplied parameters and zero_grad clears only those parameters’ accumulators. [Sources] source:pytorch-optimization-tutorial.

step() uses the gradients to change the weights

print("Before: w, v =", round(w.item(), 2), round(v.item(), 2))
print("Before: loss =", round(loss.item(), 4))
optimizer.step()
with torch.no_grad():
    new_loss = objective()
print("After:  w, v =", round(w.item(), 2), round(v.item(), 2))
print("After:  loss =", round(new_loss.item(), 4))

Output:

Before: w, v = 1.0 1.0
Before: loss = 0.5
After:  w, v = 0.8 0.8
After:  loss = 0.0392

Here \(w,v: 1\to 1-0.1(2)=0.8\), and loss decreases from \(0.5\) to \(0.0392\).

no_grad() skips recording this evaluation; the next training forward pass records again.

About three minutes. The printed loss is recomputed after the step; the old loss tensor still contains the old forward value. Ordinary optimizers update parameters without tracking the update as part of the next differentiation graph. One successful step in this example does not imply every step decreases loss for any learning rate. The optimizer does not automatically clear .grad after stepping. [Sources] source:pytorch-optimization-tutorial; source:pytorch-autograd-2-7.

The same operations train the product-image MLP

Implementation · From two scalar weights to 50,890 parameters

Data batch → model → scalar loss → backward → optimizer step.

Now assemble the complete loop, including validation.

30-second orientation. The differentiation mechanism stays the same; tensor operations replace scalar operations. Use the existing Fashion-MNIST train/validation indices and raw archived pixels. This is a PyTorch implementation demonstration, not a claimed reproduction of the earlier scikit-learn optimizer experiment.

A dataset provides examples; a loader assembles batches

from torch.utils.data import DataLoader, TensorDataset
# X_train: float32 [12000, 1, 28, 28], pixels in [0, 1]
# y_train: int64 [12000], class indices 0,...,9
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(
    train_dataset, batch_size=64, shuffle=True)
val_loader = DataLoader(
    TensorDataset(X_val, y_val), batch_size=256)
xb, yb = next(iter(train_loader))
print(xb.shape, yb.shape)

Output:

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

Dataset: train_dataset[i] returns one image–label pair; len(train_dataset) is the number of examples.

DataLoader: iterates over batches of those pairs, handling batch size and shuffling. Here next(iter(train_loader)) returns 64 examples.

About three minutes. The companion script supplies X_train, y_train, X_val, y_val from the archived dataset and existing fixed split; the slide is intentionally about the loader, not binary file parsing. TensorDataset pairs tensors by their first index; it is one implementation of the Dataset interface. The dataset defines access to individual examples, while the loader controls how to group and iterate over them. One epoch visits the training set once. Four image axes are batch, channel, height, width. Inputs need not require gradients. Shuffle training examples, preserve the held-out validation split, and keep test data outside this loop. The last batch can be smaller. [Sources] source:pytorch-optimization-tutorial; source:fashion-mnist.

nn.Module collects the model’s trainable parameters

from torch import nn
torch.manual_seed(570)
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 64),
    nn.ReLU(),
    nn.Linear(64, 10))
print(sum(p.numel() for p in model.parameters()))

Output:

50890

Batch shapes: [B, 1, 28, 28] → [B, 784] → [B, 64] → [B, 10].

model.parameters() collects the trainable weights and biases, already marked requires_grad=True.

About three minutes. Reconnect to the previously counted 50,240 plus 650 parameters. A linear layer is technically affine when bias is nonzero. Flatten has no trainable parameters; ReLU has no trainable parameters. Each linear layer registers Parameter objects so module traversal and optimizers can find them. Do not set requires_grad manually on every layer. [Sources] source:pytorch-optimization-tutorial.

Cross-entropy consumes scores before softmax

loss_fn = nn.CrossEntropyLoss()
logits = model(xb)       # [B, 10] unnormalized scores
loss = loss_fn(logits, yb)
print(loss.shape)

Output:

torch.Size([])

The loss combines log-softmax with the true class’s negative log probability, then averages over this batch.

Pass logits directly. Use logits.argmax(dim=1) to choose classes; softmax is only needed when displaying probabilities.

About three minutes. Here the default mean reduction with unweighted integer class targets is the familiar batch mean cross-entropy. Do not insert a Softmax layer before CrossEntropyLoss. Log-softmax is computed stably inside the loss. Argmax is for reporting predictions, not the differentiable path used to train the model. [Sources] source:pytorch-optimization-tutorial.

One batch needs four distinct training operations

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.train()                        # set training mode only
optimizer.zero_grad(set_to_none=True)  # clear
loss = loss_fn(model(xb), yb)          # forward
loss.backward()                       # differentiate
optimizer.step()                      # update

model.train() only sets training mode for layers such as dropout and batch normalization. It does not compute gradients or update weights.

Operation	Changes
zero_grad()	Stored parameter gradients
Forward and backward()	Activations, graph, then gradients
step()	Parameter values and optimizer state

About four minutes. Reuse Adam from Arc 4; do not rederive its moments. Construct the optimizer once, outside the batch and epoch loops, so its state persists. model.train sets the training-mode flag on the module and its submodules. Special layers such as dropout and batch normalization use that flag to select their behavior. It does not run training, compute gradients, or update parameters. This Linear/ReLU MLP computes the same function in either mode. The four operations shown perform the actual update; training/evaluation modes are separate from gradient tracking. [Sources] source:pytorch-optimization-tutorial; source:pytorch-autograd-2-7.

An epoch repeats the update for every training batch

def train_one_epoch(model, loader, loss_fn, optimizer):
    model.train()  # set training mode only
    total_loss, count = 0.0, 0
    for xb, yb in loader:
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * len(yb)
        count += len(yb)
    return total_loss / count

Multiply each batch’s mean loss by its size so a short final batch receives the correct weight.

This records the losses observed during training, as the weights change.

About four minutes including reading the loop aloud. The inputs remain CPU tensors in this example. .item gives a Python number for logging without retaining a graph. The epoch loss aggregates predictions from successive parameter states; it is not a fresh evaluation of the final model on every training example. Evaluate that final model separately when comparing train/validation curves at the same state. [Sources] source:pytorch-optimization-tutorial.

Validation measures predictions without updating the model

@torch.no_grad()
def evaluate(model, loader, loss_fn):
    model.eval()  # set evaluation mode only
    total_loss, correct, count = 0.0, 0, 0
    for xb, yb in loader:
        logits = model(xb)
        total_loss += loss_fn(logits, yb).item() * len(yb)
        correct += (logits.argmax(1) == yb).sum().item()
        count += len(yb)
    return total_loss / count, correct / count

train() / eval() only select layer behavior—important for dropout and batch normalization. Neither runs training or evaluation.

no_grad() disables graph recording. Mode and gradient tracking are separate.

About four minutes. Explain the decorator as applying no_grad throughout the function. There is no optimizer step or backward. eval is relevant for layers with training-specific behavior such as dropout; no need to teach dropout mechanics here. eval alone does not disable autograd. The next call to train_one_epoch restores train mode. Aggregate per-example metrics, including the short last batch. [Sources] source:pytorch-autograd-2-7; source:pytorch-optimization-tutorial.

The complete loop joins learning and held-out evaluation

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(5):
    train_loss = train_one_epoch(
        model, train_loader, loss_fn, optimizer)
    val_loss, val_acc = evaluate(model, val_loader, loss_fn)
    print(f"{epoch + 1}: train={train_loss:.3f}, "
          f"val={val_loss:.3f}, accuracy={val_acc:.1%}")

Output:

1: train=0.843, val=0.602, accuracy=79.6%
2: train=0.538, val=0.525, accuracy=82.2%
3: train=0.489, val=0.498, accuracy=83.1%
4: train=0.448, val=0.482, accuracy=83.6%
5: train=0.420, val=0.464, accuracy=84.1%

Model → loss → gradient → update → validation. The functions above now form one complete training loop.

Five epochs are a demonstration budget; validation guides later choices. Keep the test set untouched.

About four minutes. Intended current-class stopping point. This loop depends only on the model, loaders, loss, and two functions already shown. The companion script includes the archived-data loader and runs the whole sequence. The displayed trace continues from the preceding one-batch demonstration, with a new optimizer. Outputs were captured by running these snippets in slide order with PyTorch 2.7.1 on CPU; numerical training results can vary by environment. No benchmark accuracy is promised and these five epochs do not reproduce the earlier scikit-learn 60-epoch results. A validation trace is evidence about this run, not proof of generalization. [Sources] source:pytorch-optimization-tutorial.

detach() keeps the values but cuts the gradient path

u = torch.tensor(1., requires_grad=True)
q = torch.tensor(1., requires_grad=True)
hidden = torch.relu(u * 2.).detach()
loss = 0.5 * (q * hidden - 1.).square()
loss.backward()
print(hidden.item(), hidden.requires_grad)
print(u.grad, q.grad.item())

Output:

2.0 False
None 2.0

The detached tensor has the same values, treated as constants for this backward pass. The loss trains q, but not u through this path.

Useful: train a new prediction layer on fixed features, or compute diagnostics from predictions without extending their training graph.

About three minutes. detach returns a tensor disconnected from the source’s recorded history; it does not change the original tensor. It shares storage, so use clone if an independent snapshot is needed. Here u represents the feature-producing parameters and q the new prediction layer. Detaching this path prevents this loss from training u; it does not globally freeze u along other paths or erase old gradients. For an entirely frozen feature extractor, computing its forward under no_grad avoids building the unused graph in the first place. The next slide shows a diagnostic use while keeping the actual training loss connected. [Sources] source:pytorch-autograd-2-7.

Store diagnostic values without keeping their training graphs

loss_history, score_history = [], []
logits = model(xb)
loss = loss_fn(logits, yb)
loss_history.append(loss.item())
score_history.append(logits.detach().cpu().clone())
print(type(loss_history[-1]).__name__)
print(score_history[-1].requires_grad)
loss.backward()

Output:

float
False

Use .item() for a scalar loss; store detached scores for later plots or evaluation metrics. The original loss still supports backward.

Storing loss or logits directly keeps references to their graphs. Logging should keep the numbers, not their training history.

About three minutes. This diagnostic snippet is independent of parameter updating; no optimizer step is performed. The saved scores could support a confidence histogram or later error analysis. A list of attached tensors can keep graph metadata and, before backward or when graphs are retained, saved activations alive. Ordinary backward normally releases saved activation buffers, so do not claim every stored loss after backward retains the full graph. detach removes the history reference; cpu moves GPU data to host memory, and clone ensures an independent snapshot even on CPU. Detached tensors still use memory: store only the examples or summaries needed. For an entirely separate evaluation pass, prefer no_grad around that pass. [Sources] source:pytorch-autograd-2-7.

Initialization is part of the training recipe

• PyTorch provides useful defaults: standard layers initialize their parameters when constructed; nn.Linear scales its initial values with the input width.
• Match the method you reproduce: follow a paper’s specified initialization, including its scale, distribution, and bias settings. Defaults may differ.
• Initialization can materially change training and the final model. Record the strategy and random seed; use multiple seeds when comparing results.

torch.nn.init provides strategies such as Xavier and Kaiming initialization. Choose for the architecture and activation; do not replace a paper’s recipe blindly.

About three minutes. Standard layers have sensible starting defaults, not universally optimal initialization. In PyTorch 2.7, Linear initializes weights and biases uniformly within plus/minus 1/sqrt(in_features). Do not imply every PyTorch layer uses one strategy, or that Linear’s default equals a ReLU-specific Kaiming initialization. Randomly different weights help break hidden-unit symmetry, while scale affects the signals passed between layers. Follow the paper and accompanying implementation if they specify initialization; investigate discrepancies. Setting a seed controls a random draw, not the distribution or scale. No new empirical benchmark is claimed here. [Sources] source:pytorch-linear-2-7; source:pytorch-initialization-2-7.

Training curves help diagnose the complete loop

Both losses remain high	Training falls, validation rises	Loss oscillates sharply
Check code, fitting, and representation	Revisit capacity, regularization, or stopping	Inspect initialization, scaling, and step size

Running training loss is cheap and useful on large datasets, but mixes changing weights. For a clearer divergence diagnosis, evaluate a fixed model on a representative training subset and validation data using the same loss.

Illustrative curves, not experiment results. Keep test data out of these choices.

About three minutes. Connect to the complete loop: train_one_epoch returns the running average, while evaluate uses the model at the end of the epoch. That is a useful monitoring convention, not a perfectly matched same-model comparison. A full extra training-set pass can be expensive; use a fixed representative training subset when diagnosing a suspicious gap. Compare at the same model state and with matching evaluation behavior and data-loss definitions. A gap can reflect distribution differences as well as overfitting; patterns suggest investigations rather than unique diagnoses. Initialization belongs among optimization checks, not as a guaranteed explanation. The figure is an original schematic. Select settings and stopping with validation; preserve test data for final assessment.

PyTorch connects a model, its gradients, and repeated updates

Stage	What we control
Define	Dataset, batches, model, initialization, scalar loss
Differentiate	Forward builds the graph; backward adds parameter gradients
Update	Clear old gradients, compute new ones, then call optimizer.step()
Evaluate	Use held-out data and log values without retaining graphs

The training loop stays useful when the architecture changes. Next, replace unrestricted pixel connections with local, shared filters.

About one minute. Wrap up the implementation topic. The update row is an ordered cycle: clear before recomputing, then forward/backward, then optimizer step. train/eval select layer behavior; they do not perform learning or turn gradient tracking on/off. The next section changes the model, not the underlying autograd or training cycle.