The nn.Module Base Class

All neural network components in PyTorch are built upon the torch.nn.Module base class. To create a custom model, you must define a class that inherits from nn.Module and implement the forward method, which dictates how input data is transformed into output.

Implementing a Convolutional Network

The following example demonstrates how to construct a basic convolutional neural network skeleton. The layers are initialized in the constructor, and the data flow is defined in the forward pass.

import torch
import torch.nn as nn
import torch.nn.functional as F

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # Define convolutional layers
        self.conv_a = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
        self.conv_b = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3)

    def forward(self, x):
        # Apply ReLU activation after first convolution
        x = F.relu(self.conv_a(x))
        return F.relu(self.conv_b(x))

Simple Transformation Module

This example illustrates a minimal implementation where the module performs a scalar addition on the input tensor.

class ScalarAdder(nn.Module):
    def __init__(self):
        super(ScalarAdder, self).__init__()

    def forward(self, input_tensor):
        # Add 5 to every element
        return input_tensor + 5

# Instantiate the model
model = ScalarAdder()

# Create a dummy input tensor
data = torch.tensor(10.0)
result = model(data)
print(result)

Convolution Operations

Convolutional layers are essential for processing grid-based data like images. The nn.Conv2d layer applies a 2D convolution across an input tensor.

Conv2d Parameters

• input: Input tensor of shape (minibatch, in_channels, iH, iW).
• weight: Filters of shape (out_channels, in_channels/groups, kH, kW).
• bias: Optional bias tensor.
• stride: The step size of the kernel as it moves over the input. Can be a single number or a tuple (sH, sW).
• padding: Implicit zero-padding added to both sides of the input. Default is 0.

Using Functional Convolution

You can apply convolution operations directly using the torch.nn.functional API. The code below demonstrates manual tensor setup for convolution with varying strides and padding.

import torch
import torch.nn.functional as F

# Define a 4x4 input matrix
input_matrix = torch.tensor([
    [1.0, 2.0, 0.0, 3.0],
    [0.0, 1.0, 2.0, 3.0],
    [1.0, 2.0, 1.0, 0.0],
    [2.0, 3.0, 1.0, 1.0]
])

# Define a 2x2 kernel
kernel_matrix = torch.tensor([
    [1.0, 0.0],
    [0.0, 1.0]
])

# Reshape to (Batch_size, Channels, Height, Width)
input_batch = input_matrix.reshape(1, 1, 4, 4)
kernel_batch = kernel_matrix.reshape(1, 1, 2, 2)

# Convolution with stride 1
result_stride_1 = F.conv2d(input_batch, kernel_batch, stride=1)
print("Output (stride=1):\n", result_stride_1)

# Convolution with stride 2
result_stride_2 = F.conv2d(input_batch, kernel_batch, stride=2)
print("Output (stride=2):\n", result_stride_2)

Applying Padding

Padding is used to control the spatial dimensions of the output feature map. The example below applies a padding of 1 to the input.

# Convolution with stride 1 and padding 1
result_padded = F.conv2d(input_batch, kernel_batch, stride=1, padding=1)
print("Output (padding=1):\n", result_padded)