Writing raw mathematical gradients from scratch is tedious and error-prone. Modern deep learning relies on high-level frameworks. Let's learn how to build neural networks using PyTorch, the primary deep learning library for researchers worldwide.
PyTorch is built around two core features: Tensors (similar to NumPy arrays but with GPU acceleration support) andAutograd (automatic differentiation that tracks mathematical calculations to automatically compute gradients).
torch.nnThe torch.nn module provides pre-packaged modules for standard neural network structures. A linear layer (fully connected layer) is defined as nn.Linear(in_features, out_features):
import torch import torch.nn as nn # Linear layer mapping 3 inputs to 2 outputs layer = nn.Linear(3, 2)
nn.SequentialTo stack layers in a clean forward pipeline, we use nn.Sequential. It chains layers together so the output of one layer automatically becomes the input of the next.
backward()When you make a forward prediction and calculate the loss, running loss.backward()automatically calculates all gradients via the backpropagation chain rule. PyTorch tracks every operation, saving you from writing tedious calculus.
A tensor created with requires_grad=True tells PyTorch to record every operation on it. Calling backward() on a scalar result then fills in .grad for each tracked tensor. For y = x² + 2x the derivative is 2x + 2, so at x = 3 the gradient is 8:
import torch x = torch.tensor(3.0, requires_grad=True) y = x ** 2 + 2 * x y.backward() print(x.grad) # tensor(8.)
nn.Linear(in_features, out_features) operates on the last dimension and treats earlier dimensions as the batch. It computes x @ W.T + b, so an input of shape (4, 3) through nn.Linear(3, 2) gives shape (4, 2). Inputs must be floating point: torch.tensor([1, 2, 3]) creates integers and will raise a dtype error in a linear layer, so write [1.0, 2.0, 3.0] or call .float().
import torch.nn as nn
model = nn.Sequential(
nn.Linear(3, 8),
nn.ReLU(),
nn.Linear(8, 2),
)
batch = torch.randn(4, 3) # 4 samples, 3 features each
print(model(batch).shape) # torch.Size([4, 2])Training repeats four steps: forward pass, loss, backward pass, parameter update. Here a single linear layer learns the line y = 2x + 1 from 20 points:
x = torch.linspace(-1, 1, 20).unsqueeze(1) # shape (20, 1)
y = 2 * x + 1
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
for step in range(200):
optimizer.zero_grad() # clear old gradients
loss = loss_fn(model(x), y) # forward pass + loss
loss.backward() # compute gradients
optimizer.step() # update weights
print(model.weight.item(), model.bias.item()) # close to 2.0 and 1.0The learned weight and bias end up close to 2 and 1. Every PyTorch training script, from this toy to a large CNN, follows the same cycle.
optimizer.zero_grad(): PyTorch adds new gradients to existing ones instead of replacing them.CrossEntropyLoss wants integer class indices.backward() on a non-scalar result without a gradient argument; reduce to a scalar loss first.with torch.no_grad():.Tensors hold the data, torch.nn supplies layers, autograd computes gradients through backward(), and an optimizer applies them. Check shapes at every step and you will avoid most early errors.
Let's simulate a standard PyTorch single-layer setup in Python!
torch and torch.nn as nn.nn.Sequential containing an nn.Linear(10, 2) layer followed by a Sigmoid activation (you can simulate this with custom output rules).(1, 10) containing ones (torch.ones(1, 10)) and print the output shape.Hint: Output shape should be torch.Size([1, 2])!