🎓 A1Learner
🔍

PyTorch Squeeze and Unsqueeze

Squeeze and unsqueeze remove or add dimensions of size 1. A size-1 dimension changes the shape but holds no extra values.

Remove the size-1 dimensions

data has shape (1, 3, 1). Two dimensions have size 1, and squeeze removes both.

data = torch.tensor([[[10], [20], [30]]])
print(data.shape)
print(data.squeeze().shape)
print(data.squeeze())

Remove one named dimension

A number after squeeze names one dimension. Only that dimension is removed.

print(data.squeeze(0).shape)
print(data.squeeze(2).shape)

Add a dimension back

unsqueeze(0) adds a size-1 dimension at the front. The result reads as a batch of one row. A batch is a group of rows processed together.

flat = torch.tensor([10, 20, 30])
print(flat.unsqueeze(0))
print(flat.unsqueeze(0).shape)

Example

import torch

data = torch.tensor([[[10],
                      [20],
                      [30]]])

print("Start shape:", data.shape)

flat = data.squeeze()
print("Squeezed shape:", flat.shape)
print("Squeezed values:")
print(flat)

step = flat.unsqueeze(0)
print("After unsqueeze(0):", step.shape)

back = step.unsqueeze(2)
print("Back shape:", back.shape)
print("Back values:")
print(back)

Key points

  • Use squeeze to remove every dimension of size 1
  • Use squeeze(2) to remove only dimension 2
  • Use unsqueeze to add a dimension of size 1 at one place
  • A batch of one row of three values has shape (1, 3)
  • Squeeze and unsqueeze change the shape, never the values

In this example, data starts with shape (1, 3, 1). It holds the values 10, 20, and 30. squeeze() removes both size-1 dimensions, and the shape becomes (3,). The squeezed values print as one short row. unsqueeze(0) then adds a size-1 dimension at the front. unsqueeze(2) adds one at position 2. The tensor ends with its first shape, (1, 3, 1), and the same values.

Note

squeeze with no number can remove more dimensions than planned, so naming the dimension is safer

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What shape does unsqueeze(0) give a row of three values?