We use cookies to enhance your experience on the site

Neural Networks

You already have two pillars of PyTorch: the tensor that computes on the GPU and autograd that works out the derivatives. What is missing is the model. You could assemble one bare-handed - keep the weights in a list and multiply matrices in a loop - but by the tenth layer you lose count, the way you lose count of a buffalo herd seen through binoculars. The

torch.nn

module hands you ready-made building blocks: layers, activation functions and loss functions.

A custom model is a class that inherits from nn.Module

In PyTorch you define a custom neural network model by inheriting from

nn.Module

and implementing the
forward()
method. Three other ideas are worth ruling out right away.

An architecture is not described through JSON files - JSON carries data, but it will not perform a single matrix multiplication. In PyTorch a network is Python code.

Nor do you build one using only lists and dictionaries. If the weights sat in a plain list,

model.parameters()

would have nothing to return, the optimizer would receive an empty collection, and training would run without a single error message, except the model would not move by even one step.

nn.Module

registers layers as learnable parameters.

Finally, there is no question of importing from an external API: such an API returns the answer of a model somebody else owns, while you are building your own, trained on your data. Here is what the recommended route looks like on a safari classifier that sorts tracked animals into three species.

1import torch
2import torch.nn as nn
3import torch.optim as optim
4
5class SafariClassifier(nn.Module):
6    def __init__(self, input_size, hidden_size, num_classes):
7        super(SafariClassifier, self).__init__()
8        self.hidden = nn.Linear(input_size, hidden_size)
9        self.relu = nn.ReLU()
10        self.output = nn.Linear(hidden_size, num_classes)
11
12    def forward(self, x):
13        x = self.hidden(x)
14        x = self.relu(x)
15        x = self.output(x)
16        return x
17
18model = SafariClassifier(input_size=10, hidden_size=64, num_classes=3)

The class has two parts. In

__init__

you declare what the network is made of, and

super().__init__()

has to appear there - without that line the layers never get registered. In

forward()

you describe where the data flows. You never call that method directly - you write

model(x)

and PyTorch runs it for you.

nn.Linear(10, 64)

is a dense layer - it takes 10 numbers, returns 64 and holds the weight matrix itself. For a simple stack of layers there is a shortcut named

nn.Sequential

, but you cannot express branching in it, so a custom model with a non-standard flow always comes back to

nn.Module

The order data travels through the network

Data moves in a fixed order: Input Layer → Hidden Layer → Activation Function (ReLU) → Output Layer.

The Input Layer is not a separate object - it is the tensor entering the model, as wide as

input_size

, here 10 measurements of a single tracked animal. The Hidden Layer mixes those numbers and produces 64 new ones. The Activation Function (ReLU) stands after it, not before, because it processes the output of the hidden layer - what comes earlier is still raw data. The Output Layer goes last, because it narrows 64 numbers down to three, one per class.

If the activation dropped out of the chain, the two linear layers would collapse into one: more weights, exactly the same expressive power.

ReLU, or max(0, x)

ReLU is a function that returns

max(0, x)

: it turns negative numbers into zero and lets positive ones through untouched. The name stands for rectified linear unit, and the entire definition really is that one comparison. The quickest way to see it is on a short tensor.

1relu = nn.ReLU()
2
3print(relu(torch.tensor([-2.0, -0.5, 0.0, 3.0])))
4# tensor([0., 0., 0., 3.])

The values -2.0 and -0.5 came out as zero, and 3.0 passed through intact. That kink in the graph at zero is non-linearity - the only reason ReLU stands in the network at all.

What ReLU is not is a data compression method: four numbers went in, four came out. It is not a data sorting function either - the order was left alone, 3.0 did not jump to the front. And it is not a clustering algorithm - clustering, KMeans for instance, groups whole observations together, while ReLU works on each single number separately.

Training: loss and the five steps of the loop

A freshly created model has random weights. To learn anything, it needs a measure of error (a loss function) and a mechanism for correcting the weights (an optimizer).

nn.CrossEntropyLoss

is the loss function for multi-class classification tasks - the ones where every observation belongs to exactly one of several disjoint classes: lion, elephant or zebra. The model returns one number per class, and the loss measures how far it landed from the true label.

It is not the choice for regression, where you predict a number such as the body mass of an animal and reach for

nn.MSELoss

instead. It does not fit clustering tasks either, because clustering has no labels, so there is nothing to compare a prediction against. And it is not a tool for text generation: all it can do is score the mistake made when picking one class out of many, it produces nothing on its own.

A trap lurks here as well: if you add

nn.Softmax()

at the end of the model and feed its output into
nn.CrossEntropyLoss
, you will not see any error - training will simply go worse, because this loss function applies softmax internally.

The loop itself has five steps, always in this order:

Forward pass - outputs = model(X)

→ Compute loss
→ optimizer.zero_grad()
→ loss.backward()
→ optimizer.step()
. Each of the five shows up in the code below with its number in a comment.

1criterion = nn.CrossEntropyLoss()
2optimizer = optim.Adam(model.parameters(), lr=0.001)
3
4for epoch in range(100):
5    for batch_X, batch_y in train_loader:
6        outputs = model(batch_X)             # 1. forward pass
7        loss = criterion(outputs, batch_y)   # 2. compute loss
8
9        optimizer.zero_grad()   # 3. clear the gradients
10        loss.backward()         # 4. compute new gradients
11        optimizer.step()        # 5. update the weights

The reasons behind this order are concrete. The forward pass comes first, because without a prediction there is nothing to compare against the truth. Then the loss - a single number saying how badly the model missed. Only once you hold that number can you call

loss.backward()

, which walks the computation back through the network and fills the

grad

field of every parameter.

The most interesting question is where

optimizer.zero_grad()

sits. PyTorch accumulates gradients by default instead of overwriting them, so without clearing them everything from the previous steps is added on top of the current one - and training will not raise a single exception, the model will just learn worse and worse. The clearing therefore has to fall before

loss.backward()

. If it landed after, it would wipe the gradients that were just computed and

optimizer.step()

would correct the weights by zero. Again with no trace in the console.

The call to

optimizer.step()

closes the cycle: it nudges every weight toward a smaller error, as far as

lr

allows. Notice, @name, that

optim.Adam

received

model.parameters()

- the very collection gathered for you by

nn.Module

That is the entire recipe. Define a model with

nn.Linear

and

nn.ReLU

layers and train it by repeating those five steps over every batch, epoch after epoch.

Saving the model weights

Training runs for hours, so the result is worth keeping. What you save is not the whole model object but

state_dict()

- a dictionary of layer names and the tensors of their weights. Pickling the entire object works too, but it ties the file to one exact class path and Python version, which is why the dictionary is the recommended route.

1torch.save(model.state_dict(), 'model.pth')
2
3model = SafariClassifier(10, 64, 3)
4model.load_state_dict(torch.load('model.pth'))
5model.eval()

Read piece by piece, that save call is

torch

save

model.state_dict()

'model.pth'

. The weights come first and the file name second - the reversed order is the most common typo here.

Loading is a two-stage affair: first you create a model with the same architecture, and only then do you pour the weights into it through

torch.load

and

load_state_dict

. The

.pth

file carries numbers alone and does not know how many layers the network has. The closing

model.eval()

switches the network into evaluation mode.

Summary

You define a custom model by inheriting from
```
nn.Module
```
and implementing
```
forward()
```
. Not through JSON files, not using only lists and dictionaries, and not by importing from an external API.

The order inside the network:

Input Layer

→

Hidden Layer

→

Activation Function (ReLU)

→

Output Layer

ReLU returns
```
max(0, x)
```
and introduces non-linearity - it is not a data compression method, not a data sorting function and not a clustering algorithm.
```
nn.Linear
```
is a dense layer, and
```
nn.Sequential
```
is only a shortcut for simple stacks.
You use
```
nn.CrossEntropyLoss
```
in multi-class classification tasks. Regression uses
```
nn.MSELoss
```
, clustering works without labels, and this function generates no text.

Training loop:

Forward pass - outputs = model(X)

→

Compute loss

→

optimizer.zero_grad()

→

loss.backward()

→

optimizer.step()

Gradients accumulate by default, which is why a missing clearing step breaks training silently.

Saving the weights:

torch.save(model.state_dict(), 'model.pth')

, loading them back:

load_state_dict

In the next lesson you will meet MLflow - a tool that remembers the parameters and metrics of every training run, so that after twenty attempts you know which one came out best. Remember one thing: a network is layers and activations arranged inside

forward()

, and learning is five loop steps repeated thousands of times.

Go to CodeWorlds→