Understanding Input/Output Shapes for CNN Layers

As you start building Convolutional Neural Networks (CNNs), one of the most common practical challenges is ensuring the output shape of one layer correctly matches the expected input shape of the next. Unlike simple linear layers where only one dimension needs consideration, convolutional and pooling layers operate on multi-dimensional grid-like data (like images), involving height, width, and channel dimensions. Understanding how these dimensions transform is fundamental for constructing valid CNN architectures.

Let's consider a typical input tensor for a 2D CNN layer (like nn.Conv2d or nn.MaxPool2d). It usually has four dimensions: $(N, C_{in}, H_{in}, W_{in})$

• $N$: Batch size (number of samples processed together).
• $C_{in}$: Number of input channels (e.g., 3 for RGB images, 1 for grayscale).
• $H_{in}$: Height of the input feature map.
• $W_{in}$: Width of the input feature map.

The batch dimension $N$ usually passes through unchanged. The main transformations happen to the channels ($C$), height ($H$), and width ($W$).

Convolutional Layers (nn.Conv2d)

The torch.nn.Conv2d layer applies a 2D convolution over an input signal composed of several input planes. The most important parameters influencing the output shape are:

• in_channels ($C_{in}$): Must match the number of channels in the input tensor.
• out_channels ($C_{out}$): Determines the number of channels produced by the convolution. This is the number of filters the layer learns.
• kernel_size: The size of the convolving kernel (filter). Can be a single int for a square kernel (e.g., 3 for 3x3) or a tuple (kH, kW) for height and width.
• stride: The step size the kernel takes as it slides across the input feature map. Defaults to 1. Can be an int or a tuple (sH, sW). A larger stride results in a smaller output feature map.
• padding: Amount of zero-padding added to the borders of the input. Defaults to 0. Can be an int or a tuple (padH, padW). Padding helps control the output spatial dimensions and can preserve border information.
• dilation: Spacing between kernel elements. Defaults to 1. Larger dilation allows the kernel to cover a wider area of the input without increasing the number of parameters (atrous convolution).

The output shape $(N, C_{out}, H_{out}, W_{out})$ is determined as follows:

1. Number of Channels ($C_{out}$): This is directly set by the out_channels parameter of the nn.Conv2d layer. Each filter produces one output channel (feature map).

2. Height ($H_{out}$) and Width ($W_{out}$): These depend on the input dimensions ($H_{in}, W_{in}$) and the layer's parameters. The formula for calculating the output height is:

   $$H_{out} = \lfloor \frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1 \rfloor$$

   And similarly for the width ($W_{out}$):

   $$W_{out} = \lfloor \frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1 \rfloor$$

   Note: If padding, dilation, kernel_size, or stride are specified as single integers, they apply to both height and width dimensions (e.g., padding[0] = padding[1] = padding). The $\lfloor \cdot \rfloor$ symbol represents the floor function (rounding down to the nearest integer).

Let's look at a common scenario with dilation = 1. The formulas simplify to:

$$H_{out} = \lfloor \frac{H_{in} + 2 \times \text{padding}[0] - \text{kernel\_size}[0]}{\text{stride}[0]} + 1 \rfloor$$ $$W_{out} = \lfloor \frac{W_{in} + 2 \times \text{padding}[1] - \text{kernel\_size}[1]}{\text{stride}[1]} + 1 \rfloor$$

Example:

Suppose we have an input tensor of shape (16, 3, 32, 32) (batch=16, channels=3, height=32, width=32). We pass it through an nn.Conv2d layer defined as:

import torch
import torch.nn as nn

conv_layer = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
# Input: N=16, Cin=3, Hin=32, Win=32
input_tensor = torch.randn(16, 3, 32, 32)

# Parameters: K=3, S=1, P=1, D=1 (default)
# H_out = floor((32 + 2*1 - 1*(3-1) - 1)/1 + 1) = floor((32 + 2 - 2 - 1)/1 + 1) = floor(31/1 + 1) = 32
# W_out = floor((32 + 2*1 - 1*(3-1) - 1)/1 + 1) = floor((32 + 2 - 2 - 1)/1 + 1) = floor(31/1 + 1) = 32
# Simplified formula (D=1):
# H_out = floor((32 + 2*1 - 3)/1 + 1) = floor(31/1 + 1) = 32
# W_out = floor((32 + 2*1 - 3)/1 + 1) = floor(31/1 + 1) = 32

output_tensor = conv_layer(input_tensor)
print(output_tensor.shape)
# Expected output: torch.Size([16, 64, 32, 32])

In this case, using kernel_size=3, stride=1, and padding=1 is a common combination that preserves the input height and width (32x32 -> 32x32), while changing the number of channels from 3 to 64. This is sometimes called "same" padding, although PyTorch doesn't have an explicit 'same' option like some other frameworks; you achieve it by setting the parameters correctly.

If we change the stride to 2 (stride=2), the output dimensions will decrease:

conv_layer_s2 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=2, padding=1)
# H_out = floor((32 + 2*1 - 3)/2 + 1) = floor(31/2 + 1) = floor(15.5 + 1) = floor(16.5) = 16
# W_out = floor((32 + 2*1 - 3)/2 + 1) = floor(31/2 + 1) = floor(15.5 + 1) = floor(16.5) = 16

output_tensor_s2 = conv_layer_s2(input_tensor)
print(output_tensor_s2.shape)
# Expected output: torch.Size([16, 64, 16, 16])

Pooling Layers (nn.MaxPool2d)

Pooling layers, such as nn.MaxPool2d, are used to reduce the spatial dimensions (downsampling) of the feature maps, making the representation more compact and robust to small translations. They operate independently on each channel.

The key parameters affecting shape are similar to nn.Conv2d, but there's no concept of out_channels because pooling doesn't change the number of channels:

• kernel_size: The size of the pooling window.
• stride: The step size of the window. Often set equal to kernel_size for non-overlapping pooling (default is kernel_size).
• padding: Amount of zero-padding added.
• dilation: Controls spacing between pooling elements.

The output shape $(N, C_{out}, H_{out}, W_{out})$ calculation follows the exact same formulas for $H_{out}$ and $W_{out}$ as the convolutional layer:

$$H_{out} = \lfloor \frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1 \rfloor$$ $$W_{out} = \lfloor \frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1 \rfloor$$

Important Difference: Pooling layers do not change the number of channels. So, $C_{out} = C_{in}$.

Example:

Let's take the output of our first conv_layer (shape=[16, 64, 32, 32]) and pass it through a common max pooling layer:

# Input from previous conv layer: N=16, Cin=64, Hin=32, Win=32
pool_layer = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) # Common setup

# Parameters: K=2, S=2, P=0, D=1 (default)
# H_out = floor((32 + 2*0 - 1*(2-1) - 1)/2 + 1) = floor((32 - 1 - 1)/2 + 1) = floor(30/2 + 1) = floor(15 + 1) = 16
# W_out = floor((32 + 2*0 - 1*(2-1) - 1)/2 + 1) = floor((32 - 1 - 1)/2 + 1) = floor(30/2 + 1) = floor(15 + 1) = 16

pooled_output = pool_layer(output_tensor)
print(pooled_output.shape)
# Expected output: torch.Size([16, 64, 16, 16])

Here, the pooling layer with a 2x2 kernel and stride 2 halves the height and width dimensions (32x32 -> 16x16), while the number of channels remains unchanged (64).

Flow of tensor dimensions through a sample Convolutional and Pooling layer sequence.

Tracking Shapes in Practice

When building complex CNNs, manually calculating shapes can become tedious and error-prone. Here are a few practical tips:

1. Print Shapes: Add print(x.shape) statements after layers during initial development to verify dimensions.
2. Use Dummy Input: Create a dummy input tensor with the expected shape and pass it through your network definition step-by-step or layer-by-layer to see how the shape evolves.
3. Helper Function: Write a small helper function that takes the layer and input shape as arguments and computes the output shape using the formulas.
4. Libraries/Tools: Some libraries or utilities (like torchinfo or pytorch-summary) can automatically summarize your model, showing output shapes for each layer given an input size.

Mastering shape calculations is a necessary step in designing and debugging CNNs. By understanding how kernel_size, stride, padding, and dilation affect the spatial dimensions, and how out_channels determines the depth, you can confidently stack layers to build effective deep learning models.