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-1to be inferred.x.flatten()gives 1-d;nn.Flattenin 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.Torx.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
viewnever copies;reshapecopies only if neededcatjoins along an existing dim,stackadds a new one
Watch out for
- Calling
.view()afterpermute/transposeand getting a stride error; use.reshape().