PyTorch Flattening Tensors
Flattening turns a tensor into one long row of values. flatten can also flatten only part of a tensor.
Flatten an image-shaped tensor
This tensor is shaped like a small image. It has three layers, each two rows tall and two values wide.
image = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]])
print(image.shape)
print(image.flatten())Keep the first dimension
Flattening joins dimensions into one row. Starting at dimension 1 keeps dimension 0 apart.
data = torch.arange(24).reshape(2, 3, 4)
print(data.flatten(start_dim=1).shape)Flatten two named dimensions
A start and an end name the run to join. Dimensions outside the run keep their shape.
print(data.flatten(0, 1).shape)
print(data.flatten(0, 1))Example
import torch
data = torch.arange(24).reshape(2, 3, 4)
print("Start shape:", data.shape)
flat = data.flatten()
print("All flat shape:", flat.shape)
print("All flat values:")
print(flat)
rows = data.flatten(start_dim=1)
print("From dim 1 shape:", rows.shape)
print("Rows:")
print(rows)Key points
- Use flatten to turn a whole tensor into one row
- Use flatten(start_dim=1) to keep the first dimension and flatten the rest
- Use flatten(0, 1) to join only dimensions 0 and 1
- The values keep their order when a tensor is flattened
- A tensor with shape (3, 2, 2) flattens into a row of 12 values
In this example, data holds 24 values with shape (2, 3, 4). flatten() with no arguments joins every dimension into one row. The flat row has shape (24,), and it prints all 24 values in order. flatten(start_dim=1) then keeps the first dimension at size 2. It joins the last two dimensions into rows of 12. The result prints as two long rows.
Note
flatten(start_dim=1) keeps one row per item in a batch, which is why it is used so often
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
What shape does flatten() give a tensor with shape (2, 3, 4)?