AI.Cruz
DEEP LEARNING

Artificial Neural Networks in PyTorch: Architecture to Training Loop

Build an ANN in PyTorch using nn.Module — define layers, pick an activation, write the training loop with the four canonical steps, and understand why zero_grad matters.

1From Perceptron to Multi-Layer Network

A single perceptron computes a weighted sum of inputs plus a bias and passes it through an activation function. The output is a single number. A multi-layer ANN stacks perceptrons: the first layer receives raw input features, each subsequent hidden layer receives the previous layer's activations, and the final output layer produces the prediction. The non-linear activation functions between layers are what give a multi-layer network its expressive power — without them, stacking layers is equivalent to a single linear transformation.

Why non-linearity matters: a linear function composed with another linear function is still linear — L₂(L₁(x)) = (W₂W₁)x + (W₂b₁ + b₂). No matter how many linear layers you stack, the network can only represent linear decision boundaries. Adding a non-linear activation (ReLU, sigmoid, tanh) after each layer breaks this — the network can now approximate any continuous function (universal approximation theorem).

2Defining an ANN with nn.Module

In PyTorch you define a network by subclassing nn.Module and implementing forward(). The layers you define in __init__ are automatically registered as parameters. The simplest architecture: input layer → hidden layer + activation → output layer.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

The network is instantiated like a regular Python class: model = Net(). You then create an optimizer: optimizer = torch.optim.SGD(model.parameters(), lr=0.01) — model.parameters() returns all tensors registered in the module tree that should be updated during training.

3The Four-Step Training Loop

1. Forward

pred = model(X)
loss = criterion(pred, y)

2. Zero grad

optimizer.zero_grad()

Clears .grad on all parameters

3. Backward

loss.backward()

Fills .grad via chain rule

4. Step

optimizer.step()

p = p - lr * p.grad for each param

4Common Pitfalls

Forgetting zero_grad
Gradients accumulate. After a few batches they overflow and training diverges. Always zero before backward.
No activation between layers
Stacked linear layers collapse to one. Always add ReLU (or similar) between hidden layers or the network can't learn non-linear patterns.
Forgetting model.train() / model.eval()
Dropout and BatchNorm behave differently in train vs. eval mode. Always set the mode before running data.

Related

Train a 2-layer ANN on XOR from scratch — manual backprop with numpy

Recall check

0 of …

Say the answer out loud before revealing — recalling is what makes it stick. ➜ You've completed the Deep Learning section!

  1. Why can't you stack linear layers without non-linear activations and get a useful deep network?

  2. What does model.parameters() return and why is it passed to the optimizer?

  3. Explain the full backward pass for a 2-layer network in words.

  4. Why is XOR a good test case for a multi-layer network but not for a single perceptron?