PyTorch nn.Conv2d

Recently, I was working on a deep learning project where I needed to implement a convolutional neural network (CNN) for image classification. The cornerstone of any CNN is the convolutional layer, and in PyTorch, this is implemented through the nn.Conv2d module.

While working with this module, I realized that mastering its parameters and understanding how it processes data are crucial for building effective neural networks.

In this article, I’ll share my experience with PyTorch’s nn.Conv2d and show you how to use it effectively in your deep learning projects.

PyTorch nn.Conv2d

The nn.Conv2d is a class in PyTorch that applies a 2D convolution over an input signal composed of several input planes. In simpler terms, it’s the building block that allows neural networks to process and extract features from images.

When working with image data (which is typically represented as 3D tensors with dimensions for height, width, and channels), nn.Conv2d helps in extracting spatial features by applying filters (also called kernels) across the image.

Basic Syntax and Parameters

Let’s first understand the basic syntax of nn.Conv2d:

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None)

Here’s what each parameter means:

  1. in_channels: Number of channels in the input image (e.g., 3 for RGB images)
  2. out_channels: Number of channels produced by the convolution
  3. kernel_size: Size of the convolving kernel
  4. stride: Stride of the convolution (default: 1)
  5. padding: Zero-padding added to both sides of the input (default: 0)
  6. dilation: Spacing between kernel elements (default: 1)
  7. groups: Number of blocked connections from input to output channels (default: 1)
  8. bias: If True, adds a learnable bias to the output (default: True)
  9. padding_mode: ‘zeros’, ‘reflect’, ‘replicate’, or ‘circular’ (default: ‘zeros’)

Read PyTorch Model Eval

Methods to Use PyTorch nn.Conv2d

Now, I will explain to you the methods to use the PyTorch nn.Conv2d.

1: Create a Basic Convolutional Layer

The simplest way to use nn.Conv2d is within a neural network model. Here’s how you can create a basic convolutional layer:

import torch
import torch.nn as nn

# Create a convolutional layer with 3 input channels, 16 output channels, and a 3x3 kernel
conv_layer = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)

# Create a sample input (batch_size, channels, height, width)
input_image = torch.randn(1, 3, 32, 32)

# Apply convolution
output = conv_layer(input_image)

print(f"Input shape: {input_image.shape}")
print(f"Output shape: {output.shape}")

When I run this code, I get:

Input shape: torch.Size([1, 3, 32, 32])
Output shape: torch.Size([1, 16, 32, 32])

You can refer to the screenshot below to see the output:

pytorch conv2d example

Notice how the output has 16 channels now (the out_channels we specified) while maintaining the same spatial dimensions due to the padding we added.

2: Create a CNN with Multiple Conv2d Layers

In real-world applications, we typically use multiple convolutional layers. Here’s how I’d create a simple CNN for classifying images from the CIFAR-10 dataset:

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

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # First convolutional layer: 3 input channels (RGB), 16 output channels, 3x3 kernel
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        # Second convolutional layer: 16 input channels, 32 output channels, 3x3 kernel
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        # Third convolutional layer: 32 input channels, 64 output channels, 3x3 kernel
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        # Fully connected layer: 64*4*4 input features, 10 output features (for 10 classes)
        self.fc = nn.Linear(64 * 4 * 4, 10)

    def forward(self, x):
        # Apply conv1, then ReLU, then max pooling
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        # Apply conv2, then ReLU, then max pooling
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        # Apply conv3, then ReLU, then max pooling
        x = F.max_pool2d(F.relu(self.conv3(x)), 2)
        # Flatten the tensor
        x = x.view(-1, 64 * 4 * 4)
        # Apply the fully connected layer
        x = self.fc(x)
        return x

# Create the model
model = SimpleCNN()

# Print the model architecture
print(model)

Output:

SimpleCNN(
  (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv3): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fc): Linear(in_features=1024, out_features=10, bias=True)
)

You can refer to the screenshot below to see the output:

pytorch conv2d

This model uses three Conv2d layers, each followed by a ReLU activation function and max pooling to reduce spatial dimensions. The final output is fed into a fully connected layer to produce class probabilities.

Check out PyTorch Early Stopping

3: Understand Stride and Padding Effects

The stride and padding parameters significantly affect the output dimensions of the convolutional layer. Let’s see how:

import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# Create an input tensor (1 batch, 1 channel, 8x8 image)
input_tensor = torch.zeros(1, 1, 8, 8)
# Create a simple pattern for visualization
input_tensor[0, 0, 2:6, 2:6] = 1.0

# Display the input tensor
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.imshow(input_tensor[0, 0].numpy(), cmap='gray')
plt.title('Input')

# Conv2d with stride=1, padding=0
conv1 = nn.Conv2d(1, 1, kernel_size=3, stride=1, padding=0)
# Initialize weights to highlight edges
conv1.weight.data = torch.tensor([[[[1.0, 0.0, -1.0],
                                    [1.0, 0.0, -1.0],
                                    [1.0, 0.0, -1.0]]]])
conv1.bias.data = torch.tensor([0.0])
output1 = conv1(input_tensor)

# Display the output with stride=1, padding=0
plt.subplot(1, 3, 2)
plt.imshow(output1[0, 0].detach().numpy(), cmap='gray')
plt.title('Stride=1, Padding=0')

# Conv2d with stride=2, padding=1
conv2 = nn.Conv2d(1, 1, kernel_size=3, stride=2, padding=1)
# Use same weights as conv1
conv2.weight.data = conv1.weight.data
conv2.bias.data = conv1.bias.data
output2 = conv2(input_tensor)

# Display the output with stride=2, padding=1
plt.subplot(1, 3, 3)
plt.imshow(output2[0, 0].detach().numpy(), cmap='gray')
plt.title('Stride=2, Padding=1')

plt.tight_layout()
plt.show()

# Print output shapes
print(f"Input shape: {input_tensor.shape}")
print(f"Output shape (stride=1, padding=0): {output1.shape}")
print(f"Output shape (stride=2, padding=1): {output2.shape}")

When running this code, I observe:

Input shape: torch.Size([1, 1, 8, 8])
Output shape (stride=1, padding=0): torch.Size([1, 1, 6, 6])
Output shape (stride=2, padding=1): torch.Size([1, 1, 4, 4])

You can refer to the screenshot below to see the output:

conv2d pytorch

This clearly shows how stride and padding affect the output dimensions:

  • With stride=1 and padding=0, the output is smaller than the input.
  • With stride=2 and padding=1, the output is even smaller due to the larger stride.

Read PyTorch MSELoss

4: Implement a Real-World Example

Let’s implement a CNN for classifying images of American landmarks using PyTorch’s nn.Conv2d:

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# Define a CNN for landmark classification
class LandmarkCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(LandmarkCNN, self).__init__()
        
        # First convolutional block
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2)
        
        # Second convolutional block
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2)
        
        # Third convolutional block
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2)
        
        # Classifier
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, num_classes)
        
    def forward(self, x):
        # First block
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        
        # Second block
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        
        # Third block
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu3(x)
        x = self.pool3(x)
        
        # Classifier
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.dropout(x)
        x = self.fc2(x)
        
        return x

# Create the model
model = LandmarkCNN(num_classes=10)  # Assuming 10 landmark classes

# Example input
sample_input = torch.randn(1, 3, 32, 32)  # Batch size 1, 3 channels, 32x32 image
output = model(sample_input)

print(f"Input shape: {sample_input.shape}")
print(f"Output shape: {output.shape}")

When I run this code, I get:

Input shape: torch.Size([1, 3, 32, 32])
Output shape: torch.Size([1, 10])

This model could be used to classify images of famous American landmarks like the Statue of Liberty, the Golden Gate Bridge, Mount Rushmore, etc.

Check out PyTorch Tensor to Numpy

5: Explore Kernel Size Effects

The kernel size is another important parameter that affects how nn.Conv2d processes images. Let’s see how different kernel sizes impact feature extraction:

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

# Create a sample image (1 batch, 1 channel, 28x28 image)
sample_image = torch.zeros(1, 1, 28, 28)
# Add some patterns to the image
sample_image[0, 0, 10:20, 10:20] = 1.0  # Square
sample_image[0, 0, 5:10, 15:25] = 0.8   # Rectangle

# Apply convolutions with different kernel sizes
kernel_sizes = [3, 5, 7]
outputs = []

for k_size in kernel_sizes:
    # Create a convolution layer with specified kernel size
    conv = nn.Conv2d(1, 1, kernel_size=k_size, padding=k_size//2)
    
    # Initialize weights to detect horizontal edges
    weight = torch.zeros(1, 1, k_size, k_size)
    weight[0, 0, :k_size//2, :] = 1.0
    weight[0, 0, k_size//2+1:, :] = -1.0
    conv.weight.data = weight
    conv.bias.data = torch.tensor([0.0])
    
    # Apply convolution
    output = conv(sample_image)
    outputs.append(output)

# Visualize the results
plt.figure(figsize=(15, 5))
plt.subplot(1, 4, 1)
plt.imshow(sample_image[0, 0].numpy(), cmap='gray')
plt.title('Original Image')

for i, (k_size, output) in enumerate(zip(kernel_sizes, outputs)):
    plt.subplot(1, 4, i+2)
    plt.imshow(output[0, 0].detach().numpy(), cmap='gray')
    plt.title(f'Kernel Size: {k_size}x{k_size}')

plt.tight_layout()
plt.show()

From this experiment, I’ve observed that:

  • Smaller kernels (3×3) detect fine details and are computationally efficient
  • Medium kernels (5×5) capture more context but require more computation
  • Larger kernels (7×7) capture broader patterns but may miss fine details and are more computationally expensive

Read PyTorch Load Model

6: Custom Weight Initialization

Sometimes we need to initialize the weights of our convolutional layers with specific values. Here’s how to do that:

import torch
import torch.nn as nn
import numpy as np

# Create a convolutional layer
conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)

# Check the default weight shape
print(f"Default weight shape: {conv.weight.shape}")  # Should be [16, 3, 3, 3]

# Initialize with custom weights
# For example, Gaussian initialization with specific mean and std
torch.nn.init.normal_(conv.weight, mean=0.0, std=0.01)

# Or initialize with custom patterns
# Create Gabor filter-like kernels
for i in range(16):
    # Create a Gabor-like pattern for each output channel
    angle = i * np.pi / 8  # Different angles
    for j in range(3):  # For each input channel
        for x in range(3):
            for y in range(3):
                # Simple Gabor-like pattern
                x_c, y_c = x - 1, y - 1  # Center coordinates
                conv.weight.data[i, j, x, y] = np.cos(x_c * np.cos(angle) + y_c * np.sin(angle))

# Visualize a few of the filters
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
for i in range(8):
    plt.subplot(2, 4, i+1)
    # Show the filter for the first input channel
    plt.imshow(conv.weight.data[i, 0].detach().numpy(), cmap='viridis')
    plt.title(f'Filter {i+1}')
    plt.axis('off')

plt.tight_layout()
plt.show()

Custom weight initialization is particularly useful when:

  • You’re implementing specific filter types (e.g., edge detection, Gabor filters)
  • You’re transferring weights from a pre-trained model
  • You’re adapting a model for a new task via transfer learning

Check out PyTorch Batch Normalization

7: Use Conv2d in Transfer Learning

Transfer learning is a powerful technique where we use a pre-trained model as a starting point. Let’s see how to use nn.Conv2d in a transfer learning scenario:

import torch
import torch.nn as nn
import torchvision.models as models

# Load a pre-trained ResNet-18 model
resnet = models.resnet18(pretrained=True)

# Freeze all parameters
for param in resnet.parameters():
    param.requires_grad = False

# Modify the final layer for our landmark classification task (10 classes)
num_features = resnet.fc.in_features
resnet.fc = nn.Linear(num_features, 10)  # Replace with our custom classifier

# Now only the parameters of the final layer are trainable
trainable_params = [p for p in resnet.parameters() if p.requires_grad]
print(f"Number of trainable parameters: {len(trainable_params)}")

# Print the structure of the first convolutional layer
print("First Conv2d layer in ResNet-18:")
print(resnet.conv1)

This approach allows us to leverage the powerful feature extractors (the Conv2d layers) that have been trained on millions of images, while customizing the classifier for our specific task.

Check out PyTorch nn Linear

Understand the Math Behind Conv2d

To truly master nn.Conv2d, it helps to understand the mathematical operation it performs. A 2D convolution is essentially a sliding dot product operation:

import torch
import numpy as np
import matplotlib.pyplot as plt

# Create a simple 5x5 input
input_data = torch.tensor([
    [0, 0, 0, 0, 0],
    [0, 1, 1, 1, 0],
    [0, 1, 1, 1, 0],
    [0, 1, 1, 1, 0],
    [0, 0, 0, 0, 0]
], dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # Add batch and channel dimensions

# Define a 3x3 kernel for edge detection
kernel = torch.tensor([
    [-1, -1, -1],
    [-1,  8, -1],
    [-1, -1, -1]
], dtype=torch.float32).unsqueeze(0).unsqueeze(0)  # Add out_channels and in_channels dimensions

# Create a Conv2d layer with this kernel
conv = nn.Conv2d(1, 1, kernel_size=3, padding=0, bias=False)
conv.weight.data = kernel

# Apply the convolution
output = conv(input_data)

# Visualize input and output
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(input_data[0, 0].numpy(), cmap='gray')
plt.title('Input')
plt.subplot(1, 2, 2)
plt.imshow(output[0, 0].detach().numpy(), cmap='gray')
plt.title('Output (Edge Detection)')
plt.tight_layout()
plt.show()

print(f"Input shape: {input_data.shape}")
print(f"Output shape: {output.shape}")

The output shows how the convolution operation has detected the edges in our simple input pattern. This same principle applies whether we’re working with a simple pattern or complex real-world images.

Read Adam Optimizer PyTorch

Performance Tips for Conv2d

When working with deep networks containing many Conv2d layers, performance can become a concern. Here are some tips I’ve found helpful:

  1. Use appropriate stride values: A stride of 2 can reduce spatial dimensions and computational cost.
  2. Consider grouped convolutions: Setting the groups parameter can reduce computation.
# Standard convolution
standard_conv = nn.Conv2d(64, 128, kernel_size=3, padding=1)

# Grouped convolution (2 groups)
grouped_conv = nn.Conv2d(64, 128, kernel_size=3, padding=1, groups=2)

# Depthwise convolution (groups = in_channels)
depthwise_conv = nn.Conv2d(64, 64, kernel_size=3, padding=1, groups=64)

# Calculate parameter count
standard_params = sum(p.numel() for p in standard_conv.parameters())
grouped_params = sum(p.numel() for p in grouped_conv.parameters())
depthwise_params = sum(p.numel() for p in depthwise_conv.parameters())

print(f"Standard Conv2d parameters: {standard_params}")
print(f"Grouped Conv2d parameters: {grouped_params}")
print(f"Depthwise Conv2d parameters: {depthwise_params}")
  1. Use hardware acceleration: Ensure you’re leveraging GPU acceleration when available.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
input_tensor = input_tensor.to(device)

Over the years, I’ve found that proper understanding and configuration of PyTorch’s nn.Conv2d is essential for building effective convolutional neural networks. The module offers great flexibility through its various parameters, allowing us to tailor the network architecture to our specific needs.

Whether you’re building a model to classify American landmarks, detect objects in images, or segment medical scans, mastering nn.Conv2d will serve as a solid foundation for your deep learning journey. I hope this guide helps you understand and implement convolutional layers more effectively in your PyTorch projects.

You may also like to read:

51 Python Programs

51 PYTHON PROGRAMS PDF FREE

Download a FREE PDF (112 Pages) Containing 51 Useful Python Programs.

pyython developer roadmap

Aspiring to be a Python developer?

Download a FREE PDF on how to become a Python developer.

Let’s be friends

Be the first to know about sales and special discounts.