Tensors, Shapes and Data Batches

Objective: read the dimensions of a tensor and prevent most shape errors. Prerequisite: basic NumPy.

A tensor is a homogeneous multidimensional array. A scalar has zero axes, a vector one axis, a matrix two axes. Rank is not the number of elements: it is the number of axes.

Data	Usual shape
Dense table	`(batch, features)`
PyTorch images	`(batch, channels, height, width)`
Tokenized text	`(batch, sequence_length)`
Multiclass logits	`(batch, classes)`

python

import torch

x = torch.randn(32, 3, 224, 224)
print(x.shape, x.dtype, x.device)
assert x.ndim == 4

Batch and mini-batch

The batch groups several examples processed in parallel. A mini-batch of 32 images produces a less exact gradient estimate than the full dataset, but at a much lower cost. This small amount of noise can even help generalization.

Data type and device

Weights are usually float32. Labels expected by CrossEntropyLoss are torch.long integers. All data in an operation must be on the same device: CPU, CUDA or another accelerator.

python

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = x.to(device=device, dtype=torch.float32)

Shape discipline

At every layer, write the expected shape before the code. For a dense layer Linear(128, 10), the last incoming axis must be 128. Use reshape carefully and prefer flatten(start_dim=1) to preserve the batch axis.

Pitfalls

Confusing NHWC and NCHW.
Accidentally removing the batch axis with squeeze().
Sending the model to GPU while leaving the data on CPU.
Passing float32 labels to a loss that expects class indices.

Quick check

What is the shape of 64 RGB 32×32 images in PyTorch? Which axis does a dense layer interpret as the input dimension?

Answers

(64, 3, 32, 32). The dense layer uses the last axis.

Mastery activity — Shape audit

Start from X.shape = (32, 3, 64, 64). Feed it through a convolution producing 16 channels, a spatial reduction by two, a global pooling and a 7-class head. Write every shape and three Python assertions. Success: the batch axis remains 32 and the output is (32, 7).

Shape contract for an image batch

Reading: the B axis flows through the whole network; only the representation dimensions get transformed. Before training, explicitly check X.ndim, the channel order and the match between logits.shape[0] and y.shape[0]. The main trap is a silent axis permutation: the code may run while learning a structure that has no meaning.