PyTorch: convolution shapes and preserving the batch axis

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: convolution shapes and preserving the batch axis

Concept and contract

Trace a small NCHW image through a real convolution and flatten operation.

A batched Conv2d input convention is NCHW: examples, channels, height and width. A kernel moves over spatial dimensions while combining input channels into output channels.

Output size depends on kernel, stride, padding and dilation. Naming these axes is essential before attaching a dense layer, because flattening the batch axis can silently mix separate examples.

Mechanism and worked example

The fixture creates owned three-by-three single-channel images of ones and a two-by-two kernel with weights one and no bias.

The fixture creates owned three-by-three single-channel images of ones and a two-by-two kernel with weights one and no bias. Each output value is four, and the spatial result is two by two.

flatten(start_dim=1) preserves one row per example while combining channel and spatial features. Known weights and inputs make shape and numerical behavior independently understandable.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def image_features(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch required')
    image = torch.ones(batch_size, 1, 3, 3)
    convolution = nn.Conv2d(1, 1, kernel_size=2, bias=False)
    with torch.no_grad():
        convolution.weight.fill_(1)
    output = convolution(image)
    flattened = output.flatten(start_dim=1)
    return {'shape': list(output.shape), 'features': flattened.tolist()}

Reasoning: The fixture creates owned three-by-three single-channel images of ones and a two-by-two kernel with weights one and no bias. Each output value is four, and the spatial result is two by two.

Counterexamples and common mistakes

The counterexample flattens from dimension zero, merging all examples into one vector.

The counterexample flattens from dimension zero, merging all examples into one vector. It still produces numeric values but violates the per-example representation expected by later layers.

Another common error is confusing channels-last image arrays with NCHW tensors. Permuting axes is a semantic operation; reshaping without moving the intended axes does not generally fix that layout mismatch.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def image_features(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch required')
    image = torch.ones(batch_size, 1, 3, 3)
    convolution = nn.Conv2d(1, 1, kernel_size=2, bias=False)
    with torch.no_grad():
        convolution.weight.fill_(1)
    output = convolution(image)
    flattened = output.flatten(start_dim=0)
    return {'shape': list(output.shape), 'features': flattened.tolist()}

Reasoning: The counterexample flattens from dimension zero, merging all examples into one vector. It still produces numeric values but violates the per-example representation expected by later layers.

Independent practice: one bounded framework change

Add a second output channel with the same unit kernel initialization.

Add a second output channel with the same unit kernel initialization. Preserve the batch axis in the flattened representation and report the observed two-channel output shape.

This introduces a small channel-count change, not a larger vision architecture. Each example should now have eight flattened features, and both channels should retain the same expected convolution values.

Step 1: Add a second output channel while preserving per-example flattening. Step 2: Change out_channels on the convolution.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def image_features(batch_size):
    # Remove this placeholder, then adapt the retained baseline:
    # Add a second output channel while preserving per-example flattening.
    raise NotImplementedError
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch required')
    image = torch.ones(batch_size, 1, 3, 3)
    convolution = nn.Conv2d(1, 1, kernel_size=2, bias=False)
    with torch.no_grad():
        convolution.weight.fill_(1)
    output = convolution(image)
    flattened = output.flatten(start_dim=1)
    return {'shape': list(output.shape), 'features': flattened.tolist()}

Reasoning: Add a second output channel while preserving per-example flattening. Change out_channels on the convolution.

Hint 1: Keep all kernel weights initialized to one.
Hint 2: Verify eight features per example.

Solution, tests and trade-offs

Cases inspect one and two examples and reject an invalid batch count.

Cases inspect one and two examples and reject an invalid batch count. Practice checks the expanded channel shape and eight values per example.

Exact shapes catch a wrong flatten boundary even if the total number of values is plausible. The images are synthetic owned tensors; no image dataset, network download or GPU allocation is involved.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def image_features(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch required')
    image = torch.ones(batch_size, 1, 3, 3)
    convolution = nn.Conv2d(1, 2, kernel_size=2, bias=False)
    with torch.no_grad():
        convolution.weight.fill_(1)
    output = convolution(image)
    flattened = output.flatten(start_dim=1)
    return {'shape': list(output.shape), 'features': flattened.tolist()}

Reasoning: Cases inspect one and two examples and reject an invalid batch count. Practice checks the expanded channel shape and eight values per example.

Interview reasoning and next step

For a real vision task, document input normalization, image layout, label meaning and train-validation separation before increasing model size.

For a real vision task, document input normalization, image layout, label meaning and train-validation separation before increasing model size. Derive intermediate spatial sizes or inspect them with a controlled forward pass.

Compare against a simple baseline and evaluate held-out examples. A convolution running successfully is only a structural milestone, not evidence that the model learned useful visual features.

Optional stretch — after the core task works: Derive output sizes for explicit stride and padding changes.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-convolution. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Trace a small NCHW image through a real convolution and flatten operation.