Neural Networks in PyTorch

Neural Networks in PyTorch: nn.Module, forward(), backward(), Loss, ReLU, MaxPool, Flatten and More

If you are starting deep learning with PyTorch, you meet a handful of terms again and again: neural network, nn.Module, forward(), backward(), loss function, ReLU, MaxPool, Flatten. This post explains each one in plain language, with small code examples, and ends by putting everything together in a working image classifier.

Contents

1. What is a Neural Network?
2. nn.Module: the building block
3. forward(): the forward pass
4. Common layers: Linear, Conv2d
5. ReLU and activation functions
6. MaxPool2d
7. Flatten
8. Loss functions
9. backward(): backpropagation and autograd
10. Optimizers
11. The full training loop
12. Complete CNN example
13. Cheat sheet and common mistakes

1. What is a Neural Network?

A neural network is a mathematical function with many adjustable numbers, called parameters (weights and biases), that learns to map inputs to outputs from examples. It is loosely inspired by the brain: many simple units (neurons) connected in layers, each passing a signal to the next.

A single neuron

One neuron computes a weighted sum of its inputs, adds a bias, and passes the result through an activation function:

output = activation( w1*x1 + w2*x2 + ... + wn*xn + b )

Layers

• Input layer: receives raw data (pixels of an image, numbers in a table, tokens of text).
• Hidden layers: transform the data step by step, learning useful features.
• Output layer: produces the final prediction (class scores, a price, and so on).

How a network learns

1. Forward pass: feed input through the network and get a prediction.
2. Loss: measure how wrong the prediction is.
3. Backward pass: compute how much each parameter contributed to the error (gradients).
4. Update: nudge each parameter in the direction that reduces the loss (an optimizer does this).
5. Repeat over many batches of data.

Why "deep"? A deep neural network simply has many hidden layers. Depth lets the network build complex features from simple ones, for example edges, then shapes, then objects.

2. nn.Module: the building block

torch.nn.Module is the base class for every neural network and every layer in PyTorch. A layer such as nn.Linear is a Module, and a full model you write yourself is also a Module. Modules can contain other Modules, which is how you build large models out of small parts.

What nn.Module gives you

• Parameter tracking: any nn.Parameter or sub-module assigned in __init__ is registered automatically, so model.parameters() returns all learnable tensors.
• Callable behaviour: model(x) runs forward(x) plus hooks.
• Device moves: model.to("cuda") moves every parameter at once.
• Modes: model.train() and model.eval() switch behaviour of layers like Dropout and BatchNorm.
• Saving and loading: model.state_dict() and load_state_dict().

The standard pattern

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()          # always call this first
        self.fc1 = nn.Linear(4, 8)  # layers are defined here
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(8, 3)

    def forward(self, x):           # the computation is defined here
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet()
print(model)

Remember: __init__ declares what layers exist. forward declares how data flows through them.

Inspecting parameters

for name, p in model.named_parameters():
    print(name, p.shape)

total = sum(p.numel() for p in model.parameters())
print("Total parameters:", total)

Shortcut: nn.Sequential

For a simple straight chain of layers you can skip writing a class:

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 3)
)

3. forward(): the forward pass

The forward pass is the journey of the input data through the network to produce an output. In PyTorch you implement it in the forward method.

x = torch.randn(5, 4)   # batch of 5 samples, 4 features each
out = model(x)          # calls model.forward(x) internally
print(out.shape)        # torch.Size([5, 3])

Do not call model.forward(x) directly. Call model(x). The call operator runs registered hooks and other bookkeeping around forward.

What happens during the forward pass

1. Each layer takes a tensor and returns a new tensor.
2. While computing, PyTorch silently records every operation in a computation graph.
3. That recorded graph is what backward() later walks in reverse to compute gradients.

Tensor shapes matter

Most bugs in neural networks are shape bugs. Image batches in PyTorch use the layout (N, C, H, W): batch size, channels, height, width. A batch of 32 grayscale 28x28 images is (32, 1, 28, 28).

4. Common layers

nn.Linear (fully connected layer)

Computes y = xWT + b. Every input feature connects to every output feature.

fc = nn.Linear(in_features=128, out_features=10)
# input shape (N, 128) -> output shape (N, 10)

nn.Conv2d (convolutional layer)

Slides small learnable filters (kernels) across an image to detect local patterns such as edges and textures. Because the same filter is reused over the whole image, convolutions need far fewer parameters than fully connected layers and respect the 2D structure of images.

conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1)
# input (N, 1, 28, 28) -> output (N, 16, 28, 28)

Argument	Meaning
in_channels	Channels of the input (1 for grayscale, 3 for RGB, or the previous layer's filter count)
out_channels	Number of filters, which becomes the number of output channels
kernel_size	Filter size, for example 3 means 3x3
stride	Step size of the filter (default 1)
padding	Zeros added around the border; padding=1 with a 3x3 kernel keeps height and width unchanged

Output size formula for one spatial dimension:

out = floor( (in + 2*padding - kernel_size) / stride ) + 1

5. ReLU and activation functions

Why activations are needed

Stacking linear layers without anything in between collapses into a single linear transformation, no matter how many layers you add. An activation function adds non-linearity so the network can learn complicated, curved decision boundaries.

ReLU (Rectified Linear Unit)

ReLU(x) = max(0, x)

• Negative values become 0.
• Positive values pass through unchanged.

relu = nn.ReLU()
x = torch.tensor([-2.0, -0.5, 0.0, 1.5, 3.0])
print(relu(x))   # tensor([0.0, 0.0, 0.0, 1.5, 3.0])

Why ReLU is popular:

• Very cheap to compute.
• Its gradient is 1 for positive inputs, which helps reduce the vanishing-gradient problem seen with sigmoid and tanh.
• Produces sparse activations (many exact zeros).

Drawback: the "dying ReLU" problem. A neuron that only ever receives negative inputs outputs 0 with zero gradient and may never recover. Variants such as LeakyReLU let a small slope through for negative values.

Activation	Formula / idea	Typical use
ReLU	max(0, x)	Default for hidden layers
LeakyReLU	x if x>0, else 0.01x	Avoid dying neurons
Sigmoid	1 / (1 + e-x), range (0,1)	Binary output probability
Tanh	range (-1,1)	Recurrent networks, centred outputs
Softmax	Turns scores into probabilities summing to 1	Multi-class output
GELU	Smooth ReLU-like curve	Transformers

6. MaxPool2d

Max pooling shrinks a feature map by sliding a small window over it and keeping only the largest value in each window. It has no learnable parameters.

pool = nn.MaxPool2d(kernel_size=2, stride=2)

x = torch.tensor([[[[1., 3., 2., 4.],
                    [5., 6., 7., 8.],
                    [3., 2., 1., 0.],
                    [1., 2., 3., 4.]]]])   # shape (1, 1, 4, 4)

print(pool(x))
# tensor([[[[6., 8.],
#           [3., 4.]]]])                    # shape (1, 1, 2, 2)

Each 2x2 block is reduced to its maximum, so height and width are halved.

Why use pooling?

• Reduces size: fewer values means less computation and memory in later layers.
• Translation tolerance: a feature that shifts by a pixel or so still produces a similar pooled output.
• Keeps strong signals: the strongest activation in a region is retained.
• Helps limit overfitting by summarizing local detail.

Related layers: nn.AvgPool2d (takes the average instead of the maximum) and nn.AdaptiveAvgPool2d (you choose the output size rather than the window size).

7. Flatten

Convolution and pooling layers work on 4D tensors (N, C, H, W), but a Linear layer expects 2D input (N, features). nn.Flatten bridges them by collapsing all dimensions except the batch dimension into one.

flat = nn.Flatten()
x = torch.randn(32, 16, 7, 7)
y = flat(x)
print(y.shape)   # torch.Size([32, 784])   because 16*7*7 = 784

By default Flatten keeps dimension 0 (the batch) and flattens from dimension 1 onward. The equivalent functional form is x.view(x.size(0), -1) or torch.flatten(x, start_dim=1).

8. Loss functions

A loss function (also called cost or objective) compares the network's prediction with the true answer and returns a single number: the lower, the better. Training means minimizing this number.

Common losses

Loss	Task	Notes
nn.MSELoss	Regression	Mean of squared differences between prediction and target
nn.L1Loss	Regression	Mean absolute error, less sensitive to outliers
nn.CrossEntropyLoss	Multi-class classification	Takes raw scores (logits) and integer class labels; applies log-softmax internally
nn.BCELoss	Binary classification	Expects probabilities after a sigmoid
nn.BCEWithLogitsLoss	Binary / multi-label	Takes raw logits; more numerically stable than sigmoid + BCELoss

Example: regression

criterion = nn.MSELoss()
pred   = torch.tensor([2.5, 0.0, 2.0])
target = torch.tensor([3.0, -0.5, 2.0])
print(criterion(pred, target))   # tensor(0.1667)

Example: classification

criterion = nn.CrossEntropyLoss()
logits = torch.tensor([[2.0, 0.5, 0.1],
                       [0.2, 1.8, 0.3]])   # shape (N=2, classes=3), raw scores
labels = torch.tensor([0, 1])              # correct class index per sample
loss = criterion(logits, labels)
print(loss)

Common mistake: do not put a Softmax layer at the end of your model when using CrossEntropyLoss. The loss already applies it. Pass raw logits.

9. backward(): backpropagation and autograd

After the loss is computed, we need to know: how should each weight change to reduce the loss? The answer is the gradient of the loss with respect to each parameter. loss.backward() computes all of these gradients automatically.

The idea: chain rule on the computation graph

During the forward pass PyTorch's autograd engine records each operation. Calling backward() walks that graph from the loss back to the inputs, applying the calculus chain rule at each step. This is backpropagation.

A tiny example

w = torch.tensor(2.0, requires_grad=True)
x = torch.tensor(3.0)
y = w * x          # y = 6
loss = (y - 10) ** 2   # (6-10)^2 = 16

loss.backward()
print(w.grad)      # d(loss)/dw = 2*(y-10)*x = 2*(-4)*3 = -24

The negative gradient tells us that increasing w would decrease the loss.

Where do gradients live?

Each parameter has a .grad attribute. After loss.backward(), p.grad is filled in for every parameter with requires_grad=True (all nn.Module parameters by default).

Important rules

• Gradients accumulate. Each backward() call adds to .grad. Therefore call optimizer.zero_grad() before every new backward pass.
• backward() is called on the loss (a scalar), not on the model. You normally do not write a backward method inside nn.Module; autograd derives it from forward.
• Graph is freed after backward by default, so a new forward pass is needed for the next step.
• Use with torch.no_grad(): during evaluation to skip graph building and save memory.

Custom backward: only in advanced cases (custom torch.autograd.Function) do you write your own forward and backward static methods to define a non-standard gradient.

10. Optimizers

The gradient says which direction reduces the loss. The optimizer uses it to update the weights.

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
# or the popular adaptive choice:
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

The basic update rule of gradient descent:

weight = weight - learning_rate * gradient

• Learning rate too large: training is unstable or diverges.
• Learning rate too small: training is very slow.
• SGD is simple; Adam adapts step sizes per parameter and is a strong default.

11. The full training loop

These four calls, in this order, are the heart of PyTorch training:

for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        optimizer.zero_grad()            # 1. clear old gradients
        outputs = model(inputs)          # 2. forward pass
        loss = criterion(outputs, labels)  # 3. compute loss
        loss.backward()                  # 4. backward pass (gradients)
        optimizer.step()                 # 5. update weights

Step	Purpose
optimizer.zero_grad()	Reset gradients so they do not accumulate across batches
model(inputs)	Forward pass: produce predictions
criterion(...)	Measure the error
loss.backward()	Compute gradients via backpropagation
optimizer.step()	Apply the update to the parameters

Evaluation loop

model.eval()                      # switch to evaluation mode
correct = total = 0
with torch.no_grad():             # no gradients needed
    for inputs, labels in test_loader:
        outputs = model(inputs)
        preds = outputs.argmax(dim=1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
print("Accuracy:", correct / total)

12. Complete example: a CNN for 28x28 images

Here every concept appears together: Conv2d, ReLU, MaxPool2d, Flatten, Linear, loss, backward and optimizer. The model suits MNIST-style grayscale digit images.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),  # (N,1,28,28) -> (N,16,28,28)
            nn.ReLU(),
            nn.MaxPool2d(2),                             # -> (N,16,14,14)
            nn.Conv2d(16, 32, kernel_size=3, padding=1), # -> (N,32,14,14)
            nn.ReLU(),
            nn.MaxPool2d(2),                             # -> (N,32,7,7)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),                                # -> (N,32*7*7) = (N,1568)
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes),                 # raw logits
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

device = "cuda" if torch.cuda.is_available() else "cpu"
model = CNN().to(device)

train_data = datasets.MNIST("data", train=True, download=True,
                            transform=transforms.ToTensor())
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

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

for epoch in range(3):
    model.train()
    running = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running += loss.item()
    print(f"Epoch {epoch+1}: loss = {running/len(train_loader):.4f}")

Tracing the shapes

Layer	Output shape
Input image	(N, 1, 28, 28)
Conv2d(1,16,3,padding=1)	(N, 16, 28, 28)
ReLU	(N, 16, 28, 28)
MaxPool2d(2)	(N, 16, 14, 14)
Conv2d(16,32,3,padding=1)	(N, 32, 14, 14)
ReLU	(N, 32, 14, 14)
MaxPool2d(2)	(N, 32, 7, 7)
Flatten	(N, 1568)
Linear(1568,128) + ReLU	(N, 128)
Linear(128,10)	(N, 10) class scores

13. Cheat sheet and common mistakes

Quick reference

Term	One-line meaning
Neural network	Layered function with learnable weights that maps inputs to outputs
nn.Module	Base class for models and layers; tracks parameters
forward()	Defines how input flows through the layers
Loss function	Single number measuring prediction error
backward()	Computes gradients of the loss for every parameter via autograd
Optimizer	Updates weights using the gradients
ReLU	max(0, x); adds non-linearity
Conv2d	Learnable filters that detect local image patterns
MaxPool2d	Downsamples by keeping the maximum in each window
Flatten	Reshapes (N,C,H,W) into (N, C*H*W)
Linear	Fully connected layer: y = xWT + b

Common mistakes

1. Forgetting optimizer.zero_grad(), so gradients pile up across batches.
2. Wrong in_features for the first Linear layer after Flatten. Print the shape or compute it as channels x height x width.
3. Adding Softmax before CrossEntropyLoss. Pass raw logits.
4. Not calling model.eval() and torch.no_grad() during testing.
5. Mismatched devices: model on GPU but data on CPU (or the reverse).
6. Calling model.forward(x) instead of model(x).
7. Using the wrong label type: CrossEntropyLoss wants integer class indices (dtype long), not one-hot vectors, in the standard form.

Conclusion

A PyTorch model is an nn.Module whose forward() turns inputs into predictions. A loss function scores those predictions, backward() turns that score into gradients, and an optimizer uses the gradients to improve the weights. Layers like Conv2d, ReLU, MaxPool2d and Flatten are the pieces you arrange inside the module. Once this loop is clear, almost every deep learning model you read will feel familiar.

Happy learning! If this post helped, share it with a friend who is starting out with PyTorch.