PyTorch nn.Conv2d: Parameters, Output Shape and Examples

nn.Conv2d is the 2D convolution layer in PyTorch. You give it how many channels come in, how many should go out, and how big the kernel is, and it builds the weights for you.

conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)

It expects a 4D tensor shaped (batch, channels, height, width), and almost every problem people hit with it is a shape problem.

Versions for the runs below: PyTorch 2.14.0, Python 3.12.5. The screenshots are the terminal they came from.

nn.Conv2d output shape and parameter count

Start with the three arguments you always pass, and watch what happens to the tensor:

import torch
import torch.nn as nn

conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)

x = torch.randn(8, 3, 32, 32)          # (batch, channels, height, width)
y = conv(x)

print("input  :", tuple(x.shape))
print("output :", tuple(y.shape))
print("weight :", tuple(conv.weight.shape), " bias:", tuple(conv.bias.shape))
print("learnable parameters:", sum(p.numel() for p in conv.parameters()))

Output:

input  : (8, 3, 32, 32)
output : (8, 16, 30, 30)
weight : (16, 3, 3, 3)  bias: (16,)
learnable parameters: 448

Two things to notice. The channel count became out_channels, and the height and width shrank by 2 because a 3×3 kernel with no padding loses one pixel at each edge.

The weight is shaped (out_channels, in_channels, kernel, kernel), so the count is 16 × 3 × 3 × 3 plus one bias per output channel.

Argument	What it changes	Default
in_channels	Must match the input’s channel count	required
out_channels	The output’s channel count, and the number of filters	required
kernel_size	Window size; bigger shrinks the output more	required
stride	Step between windows; 2 roughly halves height and width	1
padding	Border pixels added before convolving	0
dilation	Gaps inside the kernel, widening its reach	1
groups	Splits channels into independent groups	1
bias	Adds one learnable offset per output channel	True

The Conv2d output size formula

You can work out the height and width before you run anything. For each spatial dimension:

out = (size + 2 * padding - dilation * (kernel - 1) - 1) // stride + 1

Rather than take that on trust, here it’s checked against what the layer actually produces, across four different settings:

import torch
import torch.nn as nn

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

x = torch.randn(1, 1, 28, 28)

for kernel, stride, padding, dilation in ((3, 1, 0, 1), (3, 1, 1, 1), (5, 2, 2, 1), (3, 1, 0, 2)):
    conv = nn.Conv2d(1, 1, kernel, stride=stride, padding=padding, dilation=dilation)
    real = tuple(conv(x).shape)[-1]
    calc = out_size(28, kernel, stride, padding, dilation)
    print(f"k={kernel} s={stride} p={padding} d={dilation} -> layer says {real:>2}, "
          f"formula says {calc:>2}, agree: {real == calc}")

Output:

k=3 s=1 p=0 d=1 -> layer says 26, formula says 26, agree: True
k=3 s=1 p=1 d=1 -> layer says 28, formula says 28, agree: True
k=5 s=2 p=2 d=1 -> layer says 14, formula says 14, agree: True
k=3 s=1 p=0 d=2 -> layer says 24, formula says 24, agree: True

The // matters. It floors, which is why an odd size with stride=2 rounds down instead of erroring.

Conv2d padding, including ‘same’ and ‘valid’

Padding adds a border of zeros before the kernel runs, which is how you stop the image shrinking at every layer:

import torch
import torch.nn as nn

x = torch.randn(1, 1, 28, 28)

for pad in (0, 1, 2, "same", "valid"):
    print(f"padding={str(pad):<8} -> {tuple(nn.Conv2d(1, 1, 3, padding=pad)(x).shape)}")

# "same" keeps height and width, but only while stride is 1
try:
    nn.Conv2d(1, 1, 3, stride=2, padding="same")
except Exception as error:
    print(f"{type(error).__name__}: {error}")

Output:

padding=0        -> (1, 1, 26, 26)
padding=1        -> (1, 1, 28, 28)
padding=2        -> (1, 1, 30, 30)
padding=same     -> (1, 1, 28, 28)
padding=valid    -> (1, 1, 26, 26)
ValueError: padding='same' is not supported for strided convolutions

padding='same' asks PyTorch to work out the border for you. It refuses when stride is anything other than 1, because the size can’t be preserved while skipping pixels.

padding='valid' is just a readable spelling of 0. The same padding rules apply to the 1D version in nn.Conv1d.

What does groups do in Conv2d?

This argument gets skipped in most tutorials, and it’s the one that makes mobile-sized models possible.

groups splits the input channels into independent sets, and each set only sees its own filters. Set it equal to in_channels and every channel is convolved on its own, which is a depthwise convolution:

import torch
import torch.nn as nn

x = torch.randn(1, 16, 32, 32)
count = lambda layer: sum(p.numel() for p in layer.parameters())

layers = (("groups=1 (plain)", nn.Conv2d(16, 16, 3, padding=1)),
          ("groups=4 (grouped)", nn.Conv2d(16, 16, 3, padding=1, groups=4)),
          ("groups=16 (depthwise)", nn.Conv2d(16, 16, 3, padding=1, groups=16)))

for name, layer in layers:
    print(f"{name:<22} weight {str(tuple(layer.weight.shape)):<16} "
          f"params {count(layer):>5}   output {tuple(layer(x).shape)}")

Output:

groups=1 (plain)       weight (16, 16, 3, 3)   params  2320   output (1, 16, 32, 32)
groups=4 (grouped)     weight (16, 4, 3, 3)    params   592   output (1, 16, 32, 32)
groups=16 (depthwise)  weight (16, 1, 3, 3)    params   160   output (1, 16, 32, 32)

The output shape never changes. What changes is the weight, whose second dimension becomes in_channels // groups.

So the weights drop by exactly that factor — 2,304 to 576 to 144 above — while the 16 bias terms stay put.

That’s the whole trick behind MobileNet and friends. A depthwise convolution followed by a 1×1 costs a fraction of a dense one, and it learns nearly as well.

nn.Conv2d or F.conv2d?

nn.Conv2d is a module that owns its weight and bias. F.conv2d is a plain function that owns nothing and expects you to pass the weight in.

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

x = torch.randn(1, 3, 8, 8)
conv = nn.Conv2d(3, 6, 3, padding=1)

from_module = conv(x)
from_function = F.conv2d(x, conv.weight, conv.bias, stride=1, padding=1)

print("identical output:", torch.allclose(from_module, from_function))
print("nn.Conv2d carries its own tensors:", [name for name, _ in conv.named_parameters()])
print("F.conv2d carries none: you hand it the weight every call")

Output:

identical output: True
nn.Conv2d carries its own tensors: ['weight', 'bias']
F.conv2d carries none: you hand it the weight every call

Use the module inside a model: it registers the parameters so the optimiser and state_dict can find them. Reach for the function when the weights come from somewhere else, as in a hypernetwork or a custom autograd experiment.

The Conv2d channel mismatch error

Nearly every Conv2d traceback you’ll see is the same one: the layer’s in_channels doesn’t match the tensor you handed it.

import torch
import torch.nn as nn

conv = nn.Conv2d(3, 16, 3)       # this layer expects 3 input channels

x = torch.randn(8, 1, 32, 32)    # but this batch has 1
conv(x)

The traceback:

Traceback (most recent call last):
  File "C:\pyguides\conv2d_channel_error.py", line 7, in <module>
    conv(x)
  File "C:\pyguides\venv\Lib\site-packages\torch\nn\modules\module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\pyguides\venv\Lib\site-packages\torch\nn\modules\module.py", line 1794, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\pyguides\venv\Lib\site-packages\torch\nn\modules\conv.py", line 565, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\pyguides\venv\Lib\site-packages\torch\nn\modules\conv.py", line 560, in _conv_forward
    return F.conv2d(
           ^^^^^^^^^
RuntimeError: Given groups=1, weight of size [16, 3, 3, 3], expected input[8, 1, 32, 32] to have 3 channels, but got 1 channels instead

Read the weight shape first. [16, 3, 3, 3] means the layer wants 3 channels, and input[8, 1, 32, 32] says you brought 1.

Greyscale images are usually what’s behind it. Either build the layer with in_channels=1, or repeat the channel three times before the forward pass.

Conv2d shapes through a small CNN

Stacking convolutions is where the shapes get away from people, so here is a small one printing its own tensor at every step:

import torch
import torch.nn as nn

class TinyCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.head = nn.Linear(32 * 8 * 8, 10)

    def forward(self, x):
        print("  input      ", tuple(x.shape))
        x = self.pool(torch.relu(self.conv1(x)))
        print("  conv1+pool ", tuple(x.shape))
        x = self.pool(torch.relu(self.conv2(x)))
        print("  conv2+pool ", tuple(x.shape))
        x = x.flatten(1)
        print("  flatten    ", tuple(x.shape))
        return self.head(x)

model = TinyCNN()
logits = model(torch.randn(4, 3, 32, 32))
print("logits:", tuple(logits.shape))
print("total parameters:", sum(p.numel() for p in model.parameters()))

Output:

  input       (4, 3, 32, 32)
  conv1+pool  (4, 16, 16, 16)
  conv2+pool  (4, 32, 8, 8)
  flatten     (4, 2048)
logits: (4, 10)
total parameters: 25578

32 × 8 × 8 in the linear layer isn’t a magic number. It’s the channel count and spatial size arriving from the layer above, which is exactly what the printed shapes show.

More PyTorch on this site:

• PyTorch nn.Conv1d
• PyTorch nn.Linear
• torch.cat for joining tensors
• torch.stack in PyTorch
• PyTorch MSELoss

Frequently asked questions