🎓 A1Learner
🔍

PyTorch Conv2d Basics

nn.Conv2d is a layer that learns kernels and slides each one over its input. Each kernel leaves one feature map in the output.

The weight tensor

Four kernels for a one-channel input give a weight shape with four rows. The bias holds one number per kernel.

conv = nn.Conv2d(1, 4, kernel_size=3)
print(conv.weight.shape)
print(conv.bias.shape)

One call, many maps

One call slides all four kernels over the image. The output holds four 3 by 3 feature maps.

conv = nn.Conv2d(1, 4, kernel_size=3)
x = torch.rand(1, 1, 5, 5)
print(conv(x).shape)

Colour input costs more

Three input channels give each kernel three slices to learn. Eight such kernels hold 224 parameters in total.

conv = nn.Conv2d(3, 8, kernel_size=3)
print(conv.weight.shape)
print(sum(p.numel() for p in conv.parameters()))

Kernels are learned

The weight tensor asks for gradients. Training can change every value inside every kernel.

print(conv.weight.requires_grad)

Example

import torch
from torch import nn

torch.manual_seed(0)

image = torch.tensor([[
    [0.0, 0.0, 1.0, 1.0, 1.0],
    [0.0, 0.0, 1.0, 1.0, 1.0],
    [0.0, 0.0, 1.0, 1.0, 1.0],
    [0.0, 0.0, 1.0, 1.0, 1.0],
    [0.0, 0.0, 1.0, 1.0, 1.0],
]]).unsqueeze(0)

kernel = torch.tensor([
    [-1.0, 0.0, 1.0],
    [-1.0, 0.0, 1.0],
    [-1.0, 0.0, 1.0],
])

by_hand = []
flat = image[0, 0]
for i in range(3):
    row = []
    for j in range(3):
        patch = flat[i:i + 3, j:j + 3]
        row.append(float((patch * kernel).sum()))
    by_hand.append(row)
by_hand = torch.tensor(by_hand)

conv = nn.Conv2d(1, 1, kernel_size=3, bias=True)
with torch.no_grad():
    conv.weight.copy_(kernel.reshape(1, 1, 3, 3))
    conv.bias.zero_()

with torch.no_grad():
    layer_out = conv(image)[0, 0]

print("By hand:")
print(by_hand)
print("Conv2d:")
print(layer_out)
print("Match:", torch.equal(by_hand, layer_out))

Key points

  • in_channels: the number of channels the layer expects in its input
  • out_channels: the number of kernels, and so the number of output maps
  • kernel_size: the side length of each square kernel
  • Use conv.weight to read or set the kernels of a layer
  • Conv2d runs the by-hand sliding sum once for each of its kernels

In this example, the seed is fixed. The image is the 5 by 5 edge image from the by-hand lesson. The program first scores the image with plain tensor maths and the edge kernel. It then builds a Conv2d layer and copies the same kernel into its weight under no_grad. Its bias is set to zero. The layer output prints next to the by-hand output, cell for cell. The last line prints the result of torch.equal on the two grids.

Note

A bias adds one learned number to each output map

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does out_channels=8 mean in nn.Conv2d?