Chapter 4: Building Neural Networks

1. nn.Module Basics

The basic building block for constructing a neural network in PyTorch is torch.nn.Module. nn.Module is a class that bundles together the functionality needed for neural networks — managing parameters (weights and biases), transferring to GPU, saving and loading models, and more. When you define your own model, you build it by inheriting from this class.

How to Create a Class that Inherits from nn.Module

To define your own model, you implement the following two methods.

• __init__(): Defines the layers used in the model
• forward(): Defines the order in which input data passes through the layers (forward propagation)

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

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

# Instantiate the model: input 10 dims → hidden layer 32 dims → output 3 dims
model = SimpleNet(input_size=10, hidden_size=32, output_size=3)
print(model)

Output:

SimpleNet(
  (fc1): Linear(in_features=10, out_features=32, bias=True)
  (relu): ReLU()
  (fc2): Linear(in_features=32, out_features=3, bias=True)
)

Running the Model and Checking Its Parameters

Once defined, the model can be called like a function to run forward propagation. Internally, model(x) calls model.forward(x).

import torch

# Dummy data with batch size 5, input dimension 10
x = torch.randn(5, 10)

# Run forward propagation
output = model(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")

# Check the parameters held by the model
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

Output:

Input shape: torch.Size([5, 10])
Output shape: torch.Size([5, 3])
fc1.weight: torch.Size([32, 10])
fc1.bias: torch.Size([32])
fc2.weight: torch.Size([3, 32])
fc2.bias: torch.Size([3])

As you can see from named_parameters(), when you assign a subclass of nn.Module (such as nn.Linear) to an attribute — as in self.fc1 = nn.Linear(...) — its parameters are automatically registered as trainable parameters of the parent model. This is one of the major benefits of using nn.Module.

2. Defining and Combining Layers

nn.Linear and nn.ReLU, used in the model above, represent a Fully Connected Layer and an Activation Function, respectively. PyTorch provides many other layers to suit various purposes.

Commonly Used Layers and Activation Functions

Layers designed for image and sequence data, such as convolutional layers (nn.Conv2d) and recurrent layers (nn.LSTM, etc.), will be covered in detail in their own dedicated series. For this chapter, let's focus on building models centered around fully connected layers.

Building Quickly with nn.Sequential

For simple models that just apply layers in sequence inside forward(), you can write more concisely using nn.Sequential.

import torch.nn as nn

model_seq = nn.Sequential(
    nn.Linear(10, 32),
    nn.ReLU(),
    nn.Linear(32, 16),
    nn.ReLU(),
    nn.Linear(16, 3)
)
print(model_seq)

Output:

Sequential(
  (0): Linear(in_features=10, out_features=32, bias=True)
  (1): ReLU()
  (2): Linear(in_features=32, out_features=16, bias=True)
  (3): ReLU()
  (4): Linear(in_features=16, out_features=3, bias=True)
)

An MLP Class that Supports a Flexible Number of Layers

In practice, you'll often want to define a model whose number and size of hidden layers can be freely changed via arguments. By combining nn.Module and nn.Sequential, you can build a general-purpose Multi-Layer Perceptron (MLP) class like the one below.

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super().__init__()
        layers = []
        prev_size = input_size
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(prev_size, hidden_size))
            layers.append(nn.ReLU())
            prev_size = hidden_size
        layers.append(nn.Linear(prev_size, output_size))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

# Input 20 dims → hidden layers [64, 32] → output 4 dims
mlp = MLP(input_size=20, hidden_sizes=[64, 32], output_size=4)
print(mlp)

Output:

MLP(
  (network): Sequential(
    (0): Linear(in_features=20, out_features=64, bias=True)
    (1): ReLU()
    (2): Linear(in_features=64, out_features=32, bias=True)
    (3): ReLU()
    (4): Linear(in_features=32, out_features=4, bias=True)
  )
)

By designing the class this way, so that it dynamically assembles layers from a list, you no longer need to rewrite the model class itself when experimenting with different hidden-layer configurations. In the second half of this chapter, we'll use this MLP class to actually train a model.

3. Loss Functions and Optimizers

Simply defining a model doesn't train anything yet. Training requires two things: a Loss Function that quantifies how wrong the model's predictions are, and an Optimizer that updates the parameters in the direction that reduces that loss.

Common Loss Functions

Common Optimizers

When in doubt, a good starting point is nn.CrossEntropyLoss() (classification) or nn.MSELoss() (regression) combined with optim.Adam(model.parameters(), lr=0.001). If training isn't stable, the common approach is to adjust the learning rate or try a different optimizer.

import torch
import torch.nn as nn
import torch.optim as optim

# Prepare a loss function and an optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(mlp.parameters(), lr=0.001)

# Compute the loss once with dummy data
x = torch.randn(8, 20)              # Batch size 8, input dimension 20
y = torch.randint(0, 4, (8,))       # Class labels 0-3 (integers)

logits = mlp(x)                     # Logits (before Softmax)
loss = criterion(logits, y)
print(f"Output shape: {logits.shape}")
print(f"Loss: {loss.item():.4f}")

Example output:

Output shape: torch.Size([8, 4])
Loss: 1.4290

As in optim.Adam(mlp.parameters(), lr=0.001), you pass the optimizer the model's full set of parameters (model.parameters()) and the learning rate. This is all that's needed so that, in the training loop we implement in the next section, simply calling optimizer.step() will automatically update these parameters.

4. Implementing the Training Loop

With the model, loss function, and optimizer all in place, it's finally time for the Training Loop. In PyTorch, the training loop for essentially any model repeats the following five steps.

1. optimizer.zero_grad(): Resets the gradients accumulated from the previous iteration
2. Forward propagation: Passes the input to the model to obtain predictions
3. Compute loss: Computes the loss from the predictions and the correct labels
4. Backward propagation (loss.backward()): Uses automatic differentiation to compute the gradient of each parameter with respect to the loss
5. optimizer.step(): Updates the parameters using the computed gradients

Let's first look at a minimal training loop that treats the entire dataset as a single batch. We'll use a synthetic dataset generated with make_classification, which will appear again in the next section.

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_classification

torch.manual_seed(42)

# Synthetic data: 4-class classification, 20 feature dimensions, 1000 samples
X, y = make_classification(
    n_samples=1000,
    n_features=20,
    n_informative=15,
    n_redundant=5,
    n_classes=4,
    random_state=42,
)
X_train = torch.tensor(X, dtype=torch.float32)
y_train = torch.tensor(y, dtype=torch.long)

model = MLP(input_size=20, hidden_sizes=[64, 32], output_size=4)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Minimal training loop (uses the entire dataset every time)
for epoch in range(5):
    optimizer.zero_grad()               # 1. Reset gradients
    outputs = model(X_train)            # 2. Forward propagation
    loss = criterion(outputs, y_train)  # 3. Compute loss
    loss.backward()                     # 4. Backward propagation
    optimizer.step()                    # 5. Update parameters

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

Example output:

Epoch 1: loss = 1.3877
Epoch 2: loss = 1.3762
Epoch 3: loss = 1.3651
Epoch 4: loss = 1.3542
Epoch 5: loss = 1.3435

This approach works, but it requires loading the entire dataset into GPU/CPU memory as one giant batch every time, which becomes impractical as the data grows. That's where Mini-batch Learning using DataLoader, covered in the next section, comes in.

5. Using DataLoader

In practice, it's common to split the dataset into small chunks (mini-batches) and run the training loop one batch at a time. In PyTorch, two classes handle this role: torch.utils.data.Dataset and torch.utils.data.DataLoader.

• Dataset: An object that, given "which sample number (index)," returns the corresponding single sample
• DataLoader: An iterator that pulls samples from a Dataset, groups them into the specified batch size, shuffles them if needed, and supplies them in order

Building Quickly with TensorDataset

If the inputs and labels are already available as Tensors, the easiest option is TensorDataset.

import torch
from torch.utils.data import TensorDataset, DataLoader

# Split into training and test data
from sklearn.model_selection import train_test_split

X_train_np, X_test_np, y_train_np, y_test_np = train_test_split(
    X, y, test_size=0.2, random_state=42
)

X_train = torch.tensor(X_train_np, dtype=torch.float32)
y_train = torch.tensor(y_train_np, dtype=torch.long)
X_test = torch.tensor(X_test_np, dtype=torch.float32)
y_test = torch.tensor(y_test_np, dtype=torch.long)

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Pull out a single batch and check its shape
batch_X, batch_y = next(iter(train_loader))
print(f"Input shape in batch: {batch_X.shape}")
print(f"Label shape in batch: {batch_y.shape}")
print(f"Dataset size: {len(train_dataset)} samples, number of batches: {len(train_loader)}")

Output:

Input shape in batch: torch.Size([32, 20])
Label shape in batch: torch.Size([32])
Dataset size: 800 samples, number of batches: 25

Setting shuffle=True shuffles the sample order at every epoch. To prevent the model from depending on the ordering of the training data, you should generally set shuffle=True for the training DataLoader (for evaluation/testing, shuffle=False is usually fine).

Creating a Custom Dataset Class

When you need more than a simple slice of a Tensor — for example, loading and preprocessing image files — you define a custom class by inheriting from Dataset. At minimum, you implement the following three methods.

import torch
from torch.utils.data import Dataset, DataLoader

class MyDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y

    def __len__(self):
        # Return the total number of samples in the dataset
        return len(self.X)

    def __getitem__(self, idx):
        # Return the sample at index idx as an (input, label) pair
        return self.X[idx], self.y[idx]

custom_dataset = MyDataset(X_train, y_train)
custom_loader = DataLoader(custom_dataset, batch_size=16, shuffle=True)

batch_X, batch_y = next(iter(custom_loader))
print(f"Custom Dataset batch shape: {batch_X.shape}, {batch_y.shape}")

Output:

Custom Dataset batch shape: torch.Size([16, 20]), torch.Size([16])

A Full Training Loop Using DataLoader

Now let's implement a more practical training loop that trains in mini-batches using DataLoader. It runs through the entire training set once per epoch and also records accuracy.

import torch.nn as nn
import torch.optim as optim

model = MLP(input_size=20, hidden_sizes=[64, 32], output_size=4)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

n_epochs = 30
for epoch in range(n_epochs):
    model.train()  # Set to training mode (enables Dropout, etc.)
    total_loss = 0.0
    correct = 0
    total = 0

    for batch_X, batch_y in train_loader:
        optimizer.zero_grad()
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * batch_X.size(0)
        preds = outputs.argmax(dim=1)
        correct += (preds == batch_y).sum().item()
        total += batch_y.size(0)

    if (epoch + 1) % 5 == 0:
        avg_loss = total_loss / total
        train_acc = correct / total
        print(f"Epoch {epoch+1:3d}/{n_epochs}: loss={avg_loss:.4f}, train_acc={train_acc:.4f}")

# Evaluate on the test data
model.eval()  # Set to evaluation mode
with torch.no_grad():
    test_outputs = model(X_test)
    test_preds = test_outputs.argmax(dim=1)
    test_acc = (test_preds == y_test).float().mean().item()
print(f"Test accuracy: {test_acc:.4f}")

Example output (execution result):

Epoch   5/30: loss=0.7789, train_acc=0.7113
Epoch  10/30: loss=0.4871, train_acc=0.8287
Epoch  15/30: loss=0.3546, train_acc=0.8788
Epoch  20/30: loss=0.2659, train_acc=0.9187
Epoch  25/30: loss=0.1975, train_acc=0.9475
Epoch  30/30: loss=0.1468, train_acc=0.9650
Test accuracy: 0.7350

While training accuracy (train_acc) has climbed to 96.5%, test accuracy remains at 73.5%. This is a classic sign of Overfitting, where the model learns the fine details of the training data too well. We'll cover countermeasures for overfitting (Dropout, regularization, early stopping, and so on) in a future series, but keep in mind that "observing both training and test accuracy" is itself an essential habit when implementing a training loop.

6. Saving and Loading Models

A model that took time to train needs to be saved so it can be reused even after the process ends. In PyTorch, the standard approach is to combine torch.save() with load_state_dict().

What is state_dict?

state_dict is a dictionary (OrderedDict) that maps the names of all parameters (weights and biases) the model holds to their values. Its key feature is that it saves and restores only the trained numerical values, not the model's structure itself.

import torch

# Save only the model's weights
torch.save(model.state_dict(), "mlp_model.pth")
print("Model saved: mlp_model.pth")

# Create a new model with the same structure and load the saved weights
loaded_model = MLP(input_size=20, hidden_sizes=[64, 32], output_size=4)
loaded_model.load_state_dict(torch.load("mlp_model.pth"))
loaded_model.eval()

# Verify that the original model and the loaded model produce the same output
with torch.no_grad():
    original_output = model(X_test)
    loaded_output = loaded_model(X_test)
    print("Outputs match:", torch.allclose(original_output, loaded_output))

Output:

Model saved: mlp_model.pth
Outputs match: True

Saving Checkpoints to Resume Training

If you want to resume training partway through, it's useful to save not just the model's weights but also the optimizer's state (optimizer state_dict), the epoch number, the loss, and so on. This kind of save is called a Checkpoint.

import torch
import torch.optim as optim

# Save multiple pieces of information together as a checkpoint
checkpoint = {
    'epoch': n_epochs,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': avg_loss,
}
torch.save(checkpoint, "checkpoint.pth")
print(f"Checkpoint saved (epoch={n_epochs}, loss={avg_loss:.4f})")

# Prepare to resume training from the checkpoint
checkpoint_loaded = torch.load("checkpoint.pth")

resumed_model = MLP(input_size=20, hidden_sizes=[64, 32], output_size=4)
resumed_model.load_state_dict(checkpoint_loaded['model_state_dict'])

resumed_optimizer = optim.Adam(resumed_model.parameters(), lr=0.001)
resumed_optimizer.load_state_dict(checkpoint_loaded['optimizer_state_dict'])

start_epoch = checkpoint_loaded['epoch']
print(f"You can resume training from epoch {start_epoch} (previous loss: {checkpoint_loaded['loss']:.4f})")

Example output:

Checkpoint saved (epoch=30, loss=0.1468)
You can resume training from epoch 30 (previous loss: 0.1468)

Restoring the optimizer's state also carries over statistics such as Adam's internal "moving averages of the gradients," letting you resume training seamlessly from where it left off. If you're only distributing a model for inference, saving the state_dict alone is sufficient — but for real-world work involving long training runs, it's worth getting comfortable with the checkpoint approach.

Exercises

import torch
import torch.nn as nn

class ThreeLayerNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 16)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(16, 8)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(8, 2)

    def forward(self, x):
        x = self.relu1(self.fc1(x))
        x = self.relu2(self.fc2(x))
        x = self.fc3(x)
        return x

model = ThreeLayerNet()
x = torch.randn(4, 8)
output = model(x)
print(f"Output shape: {output.shape}")  # torch.Size([4, 2])

import torch.nn as nn
import torch.optim as optim

# 1. Regression: output is a continuous value, so MSELoss is appropriate
regression_model = nn.Linear(13, 1)
regression_criterion = nn.MSELoss()
regression_optimizer = optim.Adam(regression_model.parameters(), lr=0.001)

# 2. 10-class classification: CrossEntropyLoss is appropriate (no Softmax on the output layer)
classification_model = nn.Linear(784, 10)
classification_criterion = nn.CrossEntropyLoss()
classification_optimizer = optim.Adam(classification_model.parameters(), lr=0.001)

print("Regression: MSELoss + Adam")
print("Classification: CrossEntropyLoss + Adam")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split

torch.manual_seed(0)

X, y = make_classification(n_samples=300, n_features=5, n_classes=2, random_state=0)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=0)

X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.long)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.long)

train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=16, shuffle=True)

class MLP(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super().__init__()
        layers = []
        prev_size = input_size
        for h in hidden_sizes:
            layers += [nn.Linear(prev_size, h), nn.ReLU()]
            prev_size = h
        layers.append(nn.Linear(prev_size, output_size))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)

model = MLP(5, [16], 2)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

for epoch in range(10):
    model.train()
    for batch_X, batch_y in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(batch_X), batch_y)
        loss.backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    test_acc = (model(X_test).argmax(dim=1) == y_test).float().mean().item()
print(f"Test accuracy: {test_acc:.4f}")

Summary

In this chapter, you learned the full workflow for building and training a neural network in PyTorch.

• ✅ You can build a model by inheriting from nn.Module, defining layers in __init__(), and defining the flow of forward propagation in forward()
• ✅ You can combine layers such as nn.Linear and nn.ReLU with nn.Sequential to build a flexible MLP
• ✅ You can choose a loss function such as nn.CrossEntropyLoss or nn.MSELoss, and an optimizer such as optim.Adam, appropriate to the type of problem (classification or regression)
• ✅ You can implement a training loop consisting of the five steps: zero_grad → forward → loss → backward → step
• ✅ You can use Dataset and DataLoader to supply data in mini-batches and train efficiently
• ✅ You understand how to save and load models via state_dict, and how to handle checkpoints that include optimizer state

Reference Resources

• PyTorch Official Documentation: torch.nn
• PyTorch Official Documentation: torch.optim
• PyTorch Official Documentation: torch.utils.data
• PyTorch Official Tutorial: Optimizing Model Parameters