Indexing and reshaping

PyTorch · Essentials · 10 min

Shape wrangling is a large share of real deep learning code, from batching images to splitting attention heads.

Indexing works like NumPy: m[0], m[:, 1], m[1:, ::2], boolean masks m[m > 8] and integer index tensors. Basic slices return views that share memory with the original; masks and index tensors return copies.

Reshaping changes how the same elements are grouped:

  • x.view(3, 4) reinterprets memory without copying; it needs a compatible memory layout.
  • x.reshape(3, 4) does the same but copies when it has to. One dimension may be -1 to be inferred.
  • x.flatten() gives 1-d; nn.Flatten in models keeps the batch dimension.
  • x.unsqueeze(0) adds a size-1 dimension (e.g. a batch of one), x.squeeze() removes them.
  • x.permute(2, 0, 1) reorders dimensions, e.g. images from H, W, C to C, H, W; x.T or x.transpose(0, 1) swaps two.

Combining: torch.cat(tensors, dim) joins along an existing dimension, torch.stack(tensors, dim) creates a new one.

Convention matters: PyTorch layers expect the batch dimension first, and images as N, C, H, W. After permute or transpose, the tensor is usually non-contiguous, so view fails with a stride error; use reshape or call .contiguous() first.

Example

import torch

x = torch.arange(12)
m = x.view(3, 4)
print(m[:, 1])                        # tensor([1, 5, 9])
print(m[m > 8])                       # tensor([ 9, 10, 11])
print(m.T.shape)                      # torch.Size([4, 3])
print(m.unsqueeze(0).shape)           # torch.Size([1, 3, 4])
print(m.reshape(-1, 6).shape)         # torch.Size([2, 6])

img = torch.zeros(32, 32, 3)          # H, W, C
print(img.permute(2, 0, 1).shape)     # torch.Size([3, 32, 32])

print(torch.cat([m, m]).shape)        # torch.Size([6, 4])
print(torch.stack([m, m]).shape)      # torch.Size([2, 3, 4])

Remember

  • Layers expect batch-first shapes; images are N, C, H, W
  • view never copies; reshape copies only if needed
  • cat joins along an existing dim, stack adds a new one

Watch out for

  • Calling .view() after permute/transpose and getting a stride error; use .reshape().

Check yourself

  1. t has shape (8, 3, 32, 32). What is torch.stack([t, t], dim=0).shape?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: pytorch.org/docs/stable/tensor_view.html