Neural Networks

Module 7: torch.nn — Building Neural Networks with PyTorch

Topics Covered

• Introduction to torch.nn
• What is nn.Module?
• Why use nn.Module?
• Anatomy of a PyTorch neural network
• Creating a custom neural network
• __init__() and forward()
• Calling a model correctly
• nn.Parameter
• Model parameters
• Model buffers
• register_buffer()
• Inspecting model architecture
• children(), modules(), and named_modules()
• parameters() and named_parameters()
• Counting trainable parameters
• Practical neural network example
• Common PyTorch model-building mistakes
• Best practices

Learning Objectives

After completing this chapter, you will be able to:

• Understand the purpose of the torch.nn package.
• Explain why nn.Module is the foundation of PyTorch models.
• Build custom neural networks using nn.Module.
• Understand the difference between __init__() and forward().
• Create and inspect trainable model parameters.
• Understand the difference between parameters and buffers.
• Register non-trainable model state using register_buffer().
• Inspect layers and model architecture programmatically.
• Calculate the number of trainable parameters.
• Follow recommended PyTorch practices when designing neural networks.

Introduction to torch.nn

Building a neural network from individual tensor operations is possible, but it quickly becomes difficult to manage as the model becomes larger.

For example, a modern neural network may contain:

• Hundreds of layers
• Millions or billions of parameters
• Different activation functions
• Normalization layers
• Dropout layers
• Embedding layers
• Attention mechanisms
• Convolution blocks
• Reusable submodules

PyTorch provides the torch.nn package to organize these components into reusable and trainable models.

pythonCopy
1import torch
2import torch.nn as nn

The torch.nn package provides many components used in deep learning, including:

• nn.Linear
• nn.Conv2d
• nn.Embedding
• nn.ReLU
• nn.GELU
• nn.Dropout
• nn.BatchNorm2d
• nn.LayerNorm
• nn.MultiheadAttention
• nn.Transformer
• Loss functions
• Model containers

The most important concept in torch.nn is nn.Module.

What Is nn.Module?

nn.Module is the base class used to construct neural network models and reusable neural network components in PyTorch.

A custom model normally inherits from nn.Module.

pythonCopy
1import torch.nn as nn
2
3class MyModel(nn.Module):
4    def __init__(self):
5        super().__init__()
6
7    def forward(self, x):
8        return x

This class does not perform any useful computation yet, but it provides the basic structure required for a PyTorch module.

Why Is nn.Module Important?

Using nn.Module provides several important capabilities.

Automatic Parameter Registration

When layers such as nn.Linear are assigned as attributes of a module, their parameters are automatically registered.

pythonCopy
1self.fc = nn.Linear(10, 5)

PyTorch can then find:

pythonCopy
1model.parameters()

and pass those parameters to an optimizer.

Device Management

A model can be moved between devices:

pythonCopy
1model.to("cuda")

or:

pythonCopy
1model.to("cpu")

Registered parameters and buffers are moved automatically.

Training and Evaluation Modes

PyTorch provides:

pythonCopy
1model.train()

and:

pythonCopy
1model.eval()

These modes are important for layers such as Dropout and Batch Normalization.

Model Serialization

Model state can be saved using:

pythonCopy
1torch.save(model.state_dict(), "model.pth")

and restored later.

Module Composition

One nn.Module can contain other modules.

For example:

textCopy
1Model
2│
3├── Linear
4├── ReLU
5├── Linear
6└── Dropout

This allows large architectures to be constructed from smaller reusable components.

Anatomy of a PyTorch Model

A typical custom PyTorch model has two important methods:

pythonCopy
1class MyModel(nn.Module):
2
3    def __init__(self):
4        super().__init__()
5
6        # Define layers here
7
8    def forward(self, x):
9        # Define data flow here
10
11        return x

The two methods have different responsibilities.

Understanding __init__()

The __init__() method is the constructor of the model.

Layers and persistent model state should generally be defined here.

Example:

pythonCopy
1import torch
2import torch.nn as nn
3
4class MyModel(nn.Module):
5
6    def __init__(self):
7        super().__init__()
8
9        self.linear = nn.Linear(10, 5)

The statement:

pythonCopy
1self.linear = nn.Linear(10, 5)

creates a Linear layer with:

textCopy
1Input features  = 10
2Output features = 5

The layer contains trainable parameters:

textCopy
1Weight: 5 × 10
2Bias:   5

Understanding forward()

The forward() method defines the computation performed by the model.

pythonCopy
1def forward(self, x):
2    return self.linear(x)

If the model receives:

pythonCopy
1x = torch.randn(4, 10)

then:

pythonCopy
1output = model(x)

produces:

textCopy
1torch.Size([4, 5])

The batch size remains 4, while the feature dimension changes from 10 to 5.

How model(x) Works

When you write:

pythonCopy
1output = model(x)

PyTorch invokes the module's call mechanism, which eventually executes the module's forward() method.

Therefore, normal PyTorch code should use:

pythonCopy
1output = model(x)

rather than directly calling:

pythonCopy
1output = model.forward(x)

Calling the module normally allows PyTorch's module machinery, including hooks and other behavior, to operate correctly.

Your First Neural Network

Let's build a simple fully connected neural network.

pythonCopy
1import torch
2import torch.nn as nn
3
4class SimpleNN(nn.Module):
5
6    def __init__(self):
7        super().__init__()
8
9        self.fc = nn.Linear(4, 2)
10
11    def forward(self, x):
12        return self.fc(x)

Create the model:

pythonCopy
1model = SimpleNN()
2
3print(model)

Output:

textCopy
1SimpleNN(
2  (fc): Linear(in_features=4, out_features=2, bias=True)
3)

Passing Data Through the Model

Create a batch containing three samples:

pythonCopy
1x = torch.randn(3, 4)
2
3output = model(x)
4
5print(output.shape)

Output:

textCopy
1torch.Size([3, 2])

The shape transformation is:

textCopy
1Input
2
3(3, 4)
4
5      │
6      ▼
7
8Linear(4, 2)
9
10      │
11      ▼
12
13Output
14
15(3, 2)

Here:

• 3 = batch size
• 4 = input features
• 2 = output features

Understanding nn.Linear

The Linear layer performs an affine transformation:

textCopy
1Y = XWᵀ + b

In PyTorch:

pythonCopy
1layer = nn.Linear(4, 2)

the weight has shape:

textCopy
1(2, 4)

and the bias has shape:

textCopy
1(2)

You can inspect them:

pythonCopy
1print(layer.weight.shape)
2print(layer.bias.shape)

Output:

textCopy
1torch.Size([2, 4])
2torch.Size([2])

Building a Multi-Layer Neural Network

A more useful network can contain multiple layers.

pythonCopy
1import torch
2import torch.nn as nn
3
4class Classifier(nn.Module):
5
6    def __init__(self):
7        super().__init__()
8
9        self.fc1 = nn.Linear(4, 16)
10        self.relu = nn.ReLU()
11        self.fc2 = nn.Linear(16, 3)
12
13    def forward(self, x):
14        x = self.fc1(x)
15        x = self.relu(x)
16        x = self.fc2(x)
17
18        return x

The architecture is:

textCopy
1Input
2  │
3  ▼
4Linear(4 → 16)
5  │
6  ▼
7ReLU
8  │
9  ▼
10Linear(16 → 3)
11  │
12  ▼
13Class Scores

Create the model:

pythonCopy
1model = Classifier()
2
3x = torch.randn(8, 4)
4
5output = model(x)
6
7print(output.shape)

Output:

textCopy
1torch.Size([8, 3])

This means the model processes eight samples and produces three output values for each sample.

Model Parameters

Trainable values such as weights and biases are represented by nn.Parameter.

For example:

pythonCopy
1layer = nn.Linear(4, 2)

contains:

textCopy
1Weight
22 × 4 = 8 parameters
3
4Bias
52 parameters
6
7Total
810 parameters

You can inspect all model parameters using:

pythonCopy
1for param in model.parameters():
2    print(param.shape)

Using named_parameters()

named_parameters() is useful when debugging or inspecting a model.

pythonCopy
1for name, param in model.named_parameters():
2    print(name, param.shape)

Example:

textCopy
1fc1.weight torch.Size([16, 4])
2fc1.bias torch.Size([16])
3fc2.weight torch.Size([3, 16])
4fc2.bias torch.Size([3])

Checking Whether Parameters Are Trainable

Every parameter has a requires_grad attribute.

pythonCopy
1for name, param in model.named_parameters():
2    print(name, param.requires_grad)

Typical output:

textCopy
1fc1.weight True
2fc1.bias True
3fc2.weight True
4fc2.bias True

A parameter with:

pythonCopy
1requires_grad=True

can receive gradients during backpropagation.

Counting Model Parameters

The numel() method returns the number of elements in a tensor.

pythonCopy
1total = sum(
2    parameter.numel()
3    for parameter in model.parameters()
4)
5
6print(total)

For the Classifier model:

textCopy
1fc1:
216 × 4 + 16 = 80
3
4fc2:
53 × 16 + 3 = 51
6
7Total:
8131

Output:

textCopy
1131

Counting Only Trainable Parameters

Sometimes a model contains frozen parameters.

Use:

pythonCopy
1trainable = sum(
2    parameter.numel()
3    for parameter in model.parameters()
4    if parameter.requires_grad
5)
6
7print(trainable)

This is particularly useful when working with:

• Transfer learning
• Fine-tuning
• Large Language Models
• LoRA
• Frozen backbone networks

What Are Model Buffers?

Not every tensor inside a neural network should be trainable.

Some tensors represent persistent state that should:

• Be saved with the model
• Move between CPU and GPU
• Not receive optimizer updates

These tensors are called buffers.

Examples include Batch Normalization running statistics.

textCopy
1Trainable parameter
2
3Weight
4Bias
5
6Non-trainable buffer
7
8Running mean
9Running variance

Registering a Buffer

PyTorch provides:

pythonCopy
1register_buffer()

Example:

pythonCopy
1import torch
2import torch.nn as nn
3
4class Example(nn.Module):
5
6    def __init__(self):
7        super().__init__()
8
9        self.register_buffer(
10            "scale",
11            torch.tensor([0.5])
12        )
13
14    def forward(self, x):
15        return x * self.scale

Now:

pythonCopy
1model = Example()
2
3print(model.scale)

Output:

textCopy
1tensor([0.5000])

Why Use register_buffer()?

Consider a tensor that should move with the model:

pythonCopy
1model.to("cuda")

A registered buffer moves with the model automatically.

A normal tensor assigned as an attribute does not receive the same parameter/buffer registration behavior.

Buffers are also included in the module's state_dict() by default.

Inspecting Buffers

Use:

pythonCopy
1for name, buffer in model.named_buffers():
2    print(name, buffer)

Output:

textCopy
1scale tensor([0.5000])

Parameters vs Buffers

The key difference is that parameters represent learnable model values, while buffers represent persistent non-learnable state.

Inspecting the Model Architecture

The simplest way to inspect a model is:

pythonCopy
1print(model)

For example:

textCopy
1Classifier(
2  (fc1): Linear(in_features=4, out_features=16, bias=True)
3  (relu): ReLU()
4  (fc2): Linear(in_features=16, out_features=3, bias=True)
5)

This provides a quick overview of the architecture.

children()

children() returns the immediate child modules.

pythonCopy
1for layer in model.children():
2    print(layer)

For the classifier:

textCopy
1Linear(in_features=4, out_features=16, bias=True)
2ReLU()
3Linear(in_features=16, out_features=3, bias=True)

This is useful when you want to inspect only the top-level components.

named_children()

You can also retrieve the names:

pythonCopy
1for name, layer in model.named_children():
2    print(name, layer)

Output:

textCopy
1fc1 Linear(...)
2relu ReLU()
3fc2 Linear(...)

modules()

modules() recursively iterates through the model and its submodules.

pythonCopy
1for module in model.modules():
2    print(type(module).__name__)

Possible output:

textCopy
1Classifier
2Linear
3ReLU
4Linear

This becomes especially useful for larger nested architectures.

named_modules()

Use named_modules() when you need both names and modules.

pythonCopy
1for name, module in model.named_modules():
2    print(name, type(module).__name__)

Example:

textCopy
1 Classifier
2fc1 Linear
3relu ReLU
4fc2 Linear

The root module has an empty name.

Model State and state_dict()

A PyTorch model's state_dict() contains the model's registered parameters and persistent buffers.

pythonCopy
1state = model.state_dict()
2
3for name, value in state.items():
4    print(name, value.shape)

For example:

textCopy
1fc1.weight torch.Size([16, 4])
2fc1.bias torch.Size([16])
3fc2.weight torch.Size([3, 16])
4fc2.bias torch.Size([3])

This is the standard mechanism used to save and restore model weights.

Saving a Model

pythonCopy
1torch.save(model.state_dict(), "model.pth")

Loading:

pythonCopy
1model = Classifier()
2
3state = torch.load(
4    "model.pth",
5    map_location="cpu"
6)
7
8model.load_state_dict(state)

For newer PyTorch workflows, it is good practice to understand the options of torch.load() and use an appropriate safe loading configuration for the type of checkpoint being loaded.

Practical Example: Binary Classification Network

Let's build a small neural network for binary classification.

pythonCopy
1import torch
2import torch.nn as nn
3
4class BinaryClassifier(nn.Module):
5
6    def __init__(self, input_features):
7        super().__init__()
8
9        self.network = nn.Sequential(
10            nn.Linear(input_features, 32),
11            nn.ReLU(),
12            nn.Linear(32, 16),
13            nn.ReLU(),
14            nn.Linear(16, 1)
15        )
16
17    def forward(self, x):
18        return self.network(x)

Create the model:

pythonCopy
1model = BinaryClassifier(10)
2
3print(model)

Create sample data:

pythonCopy
1x = torch.randn(8, 10)
2
3output = model(x)
4
5print(output.shape)

Output:

textCopy
1torch.Size([8, 1])

The architecture is:

textCopy
110 Features
2     │
3     ▼
4Linear(10 → 32)
5     │
6     ▼
7ReLU
8     │
9     ▼
10Linear(32 → 16)
11     │
12     ▼
13ReLU
14     │
15     ▼
16Linear(16 → 1)
17     │
18     ▼
19Prediction Logit

Notice that the model itself is also an nn.Module, while self.network is another module containing several child modules.

This is the foundation of hierarchical model design in PyTorch.

Parameters, Gradients, and Optimizers

A typical training pipeline looks like:

textCopy
1Input
2  │
3  ▼
4model(x)
5  │
6  ▼
7Prediction
8  │
9  ▼
10Loss
11  │
12  ▼
13loss.backward()
14  │
15  ▼
16Parameter Gradients
17  │
18  ▼
19optimizer.step()
20  │
21  ▼
22Updated Parameters

Example:

pythonCopy
1import torch
2import torch.nn as nn
3
4model = nn.Linear(4, 1)
5
6optimizer = torch.optim.SGD(
7    model.parameters(),
8    lr=0.01
9)
10
11x = torch.randn(8, 4)
12target = torch.randn(8, 1)
13
14prediction = model(x)
15
16loss = ((prediction - target) ** 2).mean()
17
18optimizer.zero_grad()
19
20loss.backward()
21
22optimizer.step()
23
24print(loss.item())

This example connects nn.Module with the Autograd and optimization concepts learned earlier.

Moving a Model to the GPU

If CUDA is available:

pythonCopy
1device = torch.device(
2    "cuda" if torch.cuda.is_available() else "cpu"
3)
4
5model = model.to(device)

Input data must also be placed on the same device:

pythonCopy
1x = x.to(device)
2
3output = model(x)

A common pattern is:

pythonCopy
1device = torch.device(
2    "cuda" if torch.cuda.is_available() else "cpu"
3)
4
5model = model.to(device)
6
7x = x.to(device)
8
9output = model(x)

This creates device-independent code that can run on either CPU or CUDA-enabled GPU systems.

Freezing Parameters

Sometimes you want to use a pretrained model without updating some of its parameters.

You can disable gradient computation for selected parameters:

pythonCopy
1for parameter in model.parameters():
2    parameter.requires_grad = False

You can then enable specific layers:

pythonCopy
1for parameter in model.fc2.parameters():
2    parameter.requires_grad = True

This technique is commonly used in transfer learning and fine-tuning.

Common Mistakes

Forgetting super().__init__()

Incorrect:

pythonCopy
1class Model(nn.Module):
2
3    def __init__(self):
4        pass

Correct:

pythonCopy
1class Model(nn.Module):
2
3    def __init__(self):
4        super().__init__()

Calling the parent constructor is essential for proper nn.Module initialization and registration behavior.

Calling forward() Directly

Avoid:

pythonCopy
1output = model.forward(x)

Prefer:

pythonCopy
1output = model(x)

Creating Trainable Layers Inside forward()

Avoid:

pythonCopy
1def forward(self, x):
2    layer = nn.Linear(10, 5)
3    return layer(x)

A new layer is created every time forward() executes, so its parameters are not persistent model parameters in the intended way.

Instead:

pythonCopy
1class Model(nn.Module):
2
3    def __init__(self):
4        super().__init__()
5
6        self.layer = nn.Linear(10, 5)
7
8    def forward(self, x):
9        return self.layer(x)

Forgetting to Register a Submodule

This is problematic:

pythonCopy
1self.layers = [
2    nn.Linear(10, 20),
3    nn.ReLU()
4]

A normal Python list does not automatically register its modules.

Use:

pythonCopy
1self.layers = nn.ModuleList([
2    nn.Linear(10, 20),
3    nn.ReLU()
4])

or:

pythonCopy
1self.layers = nn.Sequential(
2    nn.Linear(10, 20),
3    nn.ReLU()
4)

This distinction becomes important when building dynamic architectures.

Mixing CPU and GPU Tensors

This produces a device mismatch:

pythonCopy
1model = model.to("cuda")
2
3x = torch.randn(4, 10)
4
5output = model(x)

The model is on CUDA while x is on the CPU.

Correct:

pythonCopy
1device = torch.device(
2    "cuda" if torch.cuda.is_available() else "cpu"
3)
4
5model = model.to(device)
6
7x = torch.randn(4, 10).to(device)
8
9output = model(x)

Confusing Parameters with Buffers

Weights and biases that should be learned are normally parameters.

Persistent non-learnable state should generally be registered as buffers.

Best Practices

• Inherit custom neural networks from nn.Module.
• Always call super().__init__().
• Define persistent layers in __init__().
• Keep computation logic inside forward().
• Call models using model(x).
• Use named_parameters() when inspecting trainable parameters.
• Use register_buffer() for persistent non-trainable tensor state.
• Use ModuleList instead of a normal Python list for dynamically stored submodules.
• Use Sequential when the network is a straightforward chain of layers.
• Check tensor shapes during model development.
• Keep model and input tensors on the same device.
• Count trainable parameters when evaluating model size.
• Use state_dict() for saving and restoring model state.
• Freeze parameters explicitly when performing transfer learning.

Practice Exercises

Exercise 1: Create a Custom Model

Create a model with:

• Input features = 8
• Hidden features = 16
• Output features = 4

Use:

textCopy
1Linear
2ReLU
3Linear

Then print the model.

Exercise 2: Inspect Parameters

Print:

• Parameter names
• Parameter shapes
• requires_grad
• Total number of parameters
• Total number of trainable parameters

Exercise 3: Create a Buffer

Create a custom module containing:

pythonCopy
1scale = torch.tensor([0.5])

Register it using:

pythonCopy
1register_buffer()

Then print:

pythonCopy
1named_buffers()

Exercise 4: Inspect Modules

Create a model containing:

textCopy
1Linear
2ReLU
3Linear

Use:

pythonCopy
1children()

pythonCopy
1named_children()

and:

pythonCopy
1named_modules()

to inspect the architecture.

Exercise 5: GPU Compatibility

Write a program that:

1. Detects whether CUDA is available.
2. Creates a device.
3. Moves the model to the device.
4. Creates input data on the same device.
5. Performs a forward pass.

Exercise 6: Parameter Counting

Create a model with:

textCopy
1Linear(10, 32)
2Linear(32, 16)
3Linear(16, 2)

Calculate the exact total number of trainable parameters.

Exercise 7: Model State

Create a model and inspect:

pythonCopy
1model.state_dict()

Identify which entries are parameters and which entries are buffers.

Mini Project: Neural Network Inspector

Build a small utility that analyzes a PyTorch model.

Your program should display:

textCopy
1Model Name
2Number of Layers
3Total Parameters
4Trainable Parameters
5Non-Trainable Parameters
6Parameter Names
7Parameter Shapes
8Buffer Names

Example structure:

pythonCopy
1def inspect_model(model):
2
3    total = sum(
4        p.numel()
5        for p in model.parameters()
6    )
7
8    trainable = sum(
9        p.numel()
10        for p in model.parameters()
11        if p.requires_grad
12    )
13
14    print("Model:", model.__class__.__name__)
15    print("Total Parameters:", total)
16    print("Trainable Parameters:", trainable)
17
18    print("\nParameters:")
19
20    for name, parameter in model.named_parameters():
21        print(name, parameter.shape)
22
23    print("\nBuffers:")
24
25    for name, buffer in model.named_buffers():
26        print(name, buffer.shape)

This type of inspection is useful when debugging large neural networks and pretrained models.

Key Concepts to Remember

The most important ideas from this chapter are:

textCopy
1torch.nn
2   │
3   ▼
4nn.Module
5   │
6   ├── Parameters
7   │
8   ├── Buffers
9   │
10   ├── Child Modules
11   │
12   └── forward()

A PyTorch model is not simply a collection of tensor operations.

It is a structured collection of modules, parameters, buffers, and computation logic.

Understanding this structure is essential before working with CNNs, RNNs, Transformers, Vision Transformers, and Large Language Models.

Module Summary

In this chapter, you learned how torch.nn provides the foundation for building neural networks in PyTorch.

You learned:

• What torch.nn is.
• Why nn.Module is the foundation of PyTorch neural networks.
• How to create custom models.
• Why super().__init__() is required.
• How __init__() defines persistent layers.
• How forward() defines data flow.
• Why model(x) should normally be used instead of calling forward() directly.
• How nn.Linear performs an affine transformation.
• How PyTorch automatically registers parameters.
• How to inspect parameters using parameters() and named_parameters().
• How to count total and trainable parameters.
• What model buffers are.
• How to register buffers with register_buffer().
• How to inspect architectures with children(), modules(), and named_modules().
• How state_dict() stores model parameters and buffers.
• How models interact with Autograd and optimizers.
• How to move models and tensors between CPU and GPU.
• How parameter freezing is used in transfer learning.
• How to avoid common nn.Module implementation mistakes.