Shape

A tensor's shape gives dimension information.

Shape

Shape errors are the most common errors in deep learning. Layers expect specific shapes: nn.Linear wants the features in the last dimension, nn.Conv2d wants (example, channel, height, width).

x = torch.zeros(32, 3, 28, 28)
print(x.shape)          # torch.Size([32, 3, 28, 28])
print(x.shape[0])       # 32: batch size
print(x.size(1))        # 3: channels
print(len(x))           # 32

When something goes wrong, the first thing to do is print shape before and after the operation.

Adding and removing dimensions

unsqueeze adds a new axis of size 1, squeeze removes axes of size 1. Adding a batch dimension to give a single example to the model is a typical use:

x = torch.zeros(3, 28, 28)       # a single image
xb = x.unsqueeze(0)              # (1, 3, 28, 28)
print(xb.squeeze(0).shape)       # torch.Size([3, 28, 28])
v = torch.tensor([1., 2., 3.])
print(v.unsqueeze(1).shape)      # torch.Size([3, 1]): a column

Most broadcasting errors happen because of a missing unsqueeze.

Exercise Tracking shapes

For v = torch.arange(4.), print v.unsqueeze(0).shape and v.unsqueeze(1).shape on separate lines. Then print the shape of the broadcast v.unsqueeze(0) + v.unsqueeze(1) and its last row with tolist().

Show example solution

import torch

v = torch.arange(4.)

print(v.unsqueeze(0).shape)
print(v.unsqueeze(1).shape)
m = v.unsqueeze(0) + v.unsqueeze(1)
print(m.shape)
print(m[-1].tolist())

Test yourself

What is the output shape when passing a (16, 10) tensor through an <code>nn.Linear(10, 4)</code> layer?
- (16, 10)
- (4, 16)
- (16, 4)
- (10, 4)
Show answer

(16, 4) — A linear layer changes only the last dimension: 10 → 4.
Add a batch dimension to a single image.
```
xb = x.____(0)
```
Show answer
```
xb = x.unsqueeze(0)
```
Match the shapes (x: 3×28×28).

x.unsqueeze(0)

(1, 3, 28, 28)
x.shape[0]

3
x.unsqueeze(0).squeeze()

(3, 28, 28)
x.numel()

2352

Next lesson: Dtype

We Are Enhancing Your Experience

Cookie Preferences

Strictly Necessary Cookies

Analytical Cookies