PyTorch Conv Output Shapes

The conv output shape formula predicts the size of a feature map before a layer runs. It uses the input size, the kernel size, the padding, and the stride.

One worked size

Input 16, kernel 3, padding 1, and stride 1 keep the size. The arithmetic gives 16.

print((16 + 2 * 1 - 3) // 1 + 1)

Stride cuts the stops

The same input with no padding and stride 2 fits far fewer stops. The formula gives 7.

print((16 + 2 * 0 - 3) // 2 + 1)

Check against a layer

A 5 by 5 kernel with padding 2 on a 10 by 10 input should keep the size. The layer agrees.

conv = nn.Conv2d(1, 4, kernel_size=5, padding=2)
print(conv(torch.rand(1, 1, 10, 10)).shape)

Example

import torch
from torch import nn

torch.manual_seed(0)

configs = [
    (16, 3, 1, 1),
    (16, 3, 0, 2),
    (10, 5, 2, 1),
]

for size, kernel, padding, stride in configs:
    predicted_side = (size + 2 * padding - kernel) // stride + 1
    predicted = (1, 4, predicted_side, predicted_side)
    image = torch.rand(1, 1, size, size)
    conv = nn.Conv2d(1, 4, kernel_size=kernel, padding=padding, stride=stride)
    real = tuple(conv(image).shape)
    print(f"input {size}, kernel {kernel}, padding {padding}, stride {stride}")
    print("predicted:", predicted, "real:", real, "equal:", predicted == real)

Key points

Output size: (input + 2 * padding - kernel) // stride + 1
Padding adds to the input size inside the formula
Stride divides the number of stops the kernel can take
Use the formula to check a layer shape before training starts
The formula predicts the size, and a real layer run confirms it

In this example, the seed is fixed and three conv configs are tested one by one. For each config, the program first predicts the output shape with the formula. It then builds the real Conv2d layer and runs a matching random image through it. The printed lines place the predicted shape next to the real shape. Each line ends with equal: True. The formula and the layer agree on all three configs.

Note

Height and width each use the formula on their own side

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)

configs = [
    (16, 3, 1, 1),
    (16, 3, 0, 2),
    (10, 5, 2, 1),
]

for size, kernel, padding, stride in configs:
    predicted_side = (size + 2 * padding - kernel) // stride + 1
    predicted = (1, 4, predicted_side, predicted_side)
    image = torch.rand(1, 1, size, size)
    conv = nn.Conv2d(1, 4, kernel_size=kernel, padding=padding, stride=stride)
    real = tuple(conv(image).shape)
    print(f"input {size}, kernel {kernel}, padding {padding}, stride {stride}")
    print("predicted:", predicted, "real:", real, "equal:", predicted == real)

Exercise

An input of 12 meets kernel 3, padding 1, and stride 1. What is the output size?

10 11 12 14