PyTorch Layer Shapes

Layer shapes are the sizes of the tensors that pass between layers. A mismatch between two layers stops the program with a shape error.

The mismatch

The first layer returns four numbers per row. The second layer expects five, so PyTorch reports RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x4 and 5x2).

first = nn.Linear(3, 4)
second = nn.Linear(5, 2)
x = torch.ones(1, 3)
out = second(first(x))

Print shapes between layers

Printing .shape after each layer shows where the sizes stop matching. The middle shape is the clue.

first = nn.Linear(3, 4)
second = nn.Linear(4, 2)
batch = torch.tensor([[1.0, 2.0, 3.0],
                      [0.0, 1.0, 0.0]])
middle = first(batch)
print(batch.shape)
print(middle.shape)

The fix

The fix changes the second layer, not the data. Its input size now matches the middle shape.

print(second(middle).shape)

Example

import torch
from torch import nn

torch.manual_seed(0)

first = nn.Linear(3, 4)
second = nn.Linear(4, 2)

batch = torch.tensor([[1.0, 2.0, 3.0],
                      [0.0, 1.0, 0.0]])
middle = first(batch)
output = second(middle)

print("Input shape:", batch.shape)
print("After first layer:", middle.shape)
print("Output shape:", output.shape)

Key points

Use tensor.shape to print the size of a tensor at any step
Shape rule: one layer's output size must equal the next layer's input size
The liar: the layer whose input size disagrees with the shape before it
RuntimeError: the error type PyTorch raises when a shape check fails

In this example, the seed is set first, so the starting weights are repeatable. The first layer takes rows of three and returns rows of four. The second layer takes rows of four, so the shapes agree. The batch holds two rows of three numbers. The prints show the shape at each step of the pass. The middle shape matches the second layer's input size. This is the fixed version of the mismatch from the snippets. Reading shapes this way finds a wrong layer fast.

Note

The batch size can change freely, but each row must keep the same size

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn

torch.manual_seed(0)

first = nn.Linear(3, 4)
second = nn.Linear(4, 2)

batch = torch.tensor([[1.0, 2.0, 3.0],
                      [0.0, 1.0, 0.0]])
middle = first(batch)
output = second(middle)

print("Input shape:", batch.shape)
print("After first layer:", middle.shape)
print("Output shape:", output.shape)

Exercise

A tensor between two layers has shape (2, 4). What input size must the next nn.Linear have?

3 4 2 8