Convolution output shape: work through padding, stride and channels

A convolution's output shape depends on its input dimensions, kernel, stride, padding and dilation. The number of output channels is a separate parameter. Tracking those quantities explicitly helps diagnose the matrix-shape error that often appears later, when a network flattens its feature maps into a linear layer.

We will follow an invented batch of two 8-by-10 RGB images through two convolutions. No image files or pretrained weights are required. The goal is to predict dimensions before asking PyTorch, then identify exactly which assumption a mismatched layer violates.

Calculate one spatial dimension

For a standard two-dimensional convolution, apply this formula independently to height and width:

output = floor((input + 2 × padding - dilation × (kernel - 1) - 1) / stride + 1).

The effective kernel width is dilation × (kernel - 1) + 1. Dilation spreads the kernel's sampling positions; stride controls the spacing between output positions. Padding adds the border used by the operation.

For input height 8, kernel 3, padding 1, dilation 1 and stride 2, the output height is 4. For width 10 under the same settings, the output width is 5. Choosing four output channels gives a complete tensor shape of (2, 4, 4, 5) in batch-channel-height-width order.

Verify the network and reproduce the failure

The second convolution uses kernel 3, no padding and stride 1. It changes the 4-by-5 spatial dimensions to 2-by-3, keeping four output channels. Each sample therefore has 4 × 2 × 3 = 24 flattened features.

from math import floor
import torch
from torch import nn

def output_size(size, kernel, stride=1, padding=0, dilation=1):
    return floor((size + 2*padding - dilation*(kernel-1) - 1) / stride + 1)

x = torch.zeros(2, 3, 8, 10)
first = nn.Conv2d(3, 4, kernel_size=3, stride=2, padding=1)
second = nn.Conv2d(4, 4, kernel_size=3)

y = first(x)
z = second(y)
assert y.shape[-2:] == (output_size(8, 3, 2, 1), output_size(10, 3, 2, 1))
print("first:", tuple(y.shape))
print("second:", tuple(z.shape))
flat = z.flatten(start_dim=1)
print("flattened:", tuple(flat.shape))

try:
    nn.Linear(80, 5)(flat)
except RuntimeError:
    print("80-feature linear layer: shape mismatch")

correct = nn.Linear(24, 5)(flat)
print("correct output:", tuple(correct.shape))

Expected output:

first: (2, 4, 4, 5)
second: (2, 4, 2, 3)
flattened: (2, 24)
80-feature linear layer: shape mismatch
correct output: (2, 5)

The incorrect layer expected the 80 features produced by the first convolution, but receives the 24 left after the second. Changing only the batch size would not fix that contract.

Keep channels and spatial dimensions separate

Conv2d(3, 4, ...) expects three input channels and produces four output channels. It does not produce one independent output channel for each input channel and kernel position. Each output channel combines the relevant input channels according to the convolution's grouping configuration.

Likewise, flattening with start_dim=1 preserves the batch dimension. Flattening the whole tensor would mix examples into one vector, changing the meaning of the downstream computation. Shape checks should communicate what each axis represents, not merely force the total element count to fit.

The PyTorch Conv2d reference specifies the shape formula, channel relationships and grouping constraints. Our example uses ordinary groups equal to one and symmetric integer padding; other configurations require checking their actual documented rules.

Changing image size changes the contract

A linear layer with 24 input features fits this particular network and input size. Larger input images can produce a different number of flattened features. Resizing to a declared input shape is one option. Adaptive pooling is another architectural choice when a fixed output grid is appropriate.

A lazy linear layer can infer an initial feature count, but that does not make arbitrary later input shapes compatible. Nor does reshaping a tensor to the expected element count repair an incorrect semantic layout. Understand why the feature dimensions differ before changing them.

Use a shape ledger when debugging

Write the expected batch, channel, height and width after every operation. Include pooling, cropping and concatenation, not only convolution. A small dummy forward pass can verify the ledger before an expensive training run. The zeros in this script suffice for shapes; they say nothing about predictive quality.

As an exercise, change the first convolution's dilation to two while retaining its other settings. Calculate both new spatial dimensions, verify them, and update the next layer only after checking whether its kernel still fits. Then try a non-square kernel and calculate the axes separately.

Continue with PyTorch training and evaluation behavior and Deep Learning resources. Together, shape and state checks remove two common sources of confusing model failures.