Tensors & autograd in PyTorch
Everything in PyTorch — images, sentences, model weights, gradients — is a tensor: an n-dimensional array, much like a NumPy array but able to live on a GPU and to track the operations performed on it. Mastering tensors and autograd (automatic differentiation) is the foundation for every later lesson.
Creating tensors
import torch
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) # from a Python list
z = torch.zeros(3, 4) # 3x4 of zeros
r = torch.rand(2, 3) # uniform [0, 1)
n = torch.randn(2, 3) # standard normal
s = torch.arange(0, 10, 2) # tensor([0, 2, 4, 6, 8])
print(a.shape, a.dtype, a.device) # torch.Size([2, 2]) torch.float32 cpu
Ranks and shapes you will meet
| Data | Typical shape | Meaning |
|---|---|---|
| Scalar loss | () |
a single number |
| Daily sales for 30 days | (30,) |
vector |
| Batch of tabular rows | (64, 12) |
64 rows, 12 features |
| Batch of RGB images | (32, 3, 224, 224) |
N, channels, height, width |
| Batch of token IDs | (16, 128) |
16 sentences, 128 tokens each |
Most bugs in deep learning code are shape bugs, so print .shape often.
Everyday operations
x = torch.randn(64, 12)
w = torch.randn(12, 1)
y = x @ w # matrix multiply -> (64, 1)
y = y.squeeze(1) # drop size-1 dim -> (64,)
imgs = torch.rand(8, 3, 32, 32)
flat = imgs.view(8, -1) # (8, 3072); -1 lets PyTorch infer the size
mean_per_channel = imgs.mean(dim=(0, 2, 3)) # (3,)
prices = torch.tensor([120.0, 450.0, 80.0])
with_gst = prices * 1.05 # broadcasting a scalar
Broadcasting stretches dimensions of size 1 so shapes match, exactly as in NumPy. Use .reshape() when a tensor may not be contiguous in memory; .view() requires contiguity.
NumPy interop and devices
import numpy as np
arr = np.array([1.0, 2.0, 3.0])
t = torch.from_numpy(arr) # shares memory with arr
back = t.numpy()
device = "cuda" if torch.cuda.is_available() else "cpu"
t_gpu = t.to(device)
result = (t_gpu * 2).cpu() # bring back to CPU before .numpy()
Both operands of an operation must be on the same device, otherwise PyTorch raises an error.
Autograd: gradients for free
Training a network means nudging each weight in the direction that reduces the loss. That direction is the gradient. Autograd records every operation on tensors with requires_grad=True in a computation graph, then walks it backwards to compute derivatives.
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x # y = x^3 + 2x
y.backward() # dy/dx = 3x^2 + 2
print(x.grad) # tensor(14.)
A tiny linear regression shows the full idea — predicting a sweet shop's daily sales (in thousand INR) from footfall:
torch.manual_seed(0)
footfall = torch.rand(100, 1) * 10
sales = 3.0 * footfall + 5.0 + torch.randn(100, 1) * 0.5
w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
lr = 0.01
for step in range(2000):
pred = footfall * w + b
loss = ((pred - sales) ** 2).mean()
loss.backward()
with torch.no_grad(): # update without recording
w -= lr * w.grad
b -= lr * b.grad
w.grad.zero_() # gradients accumulate otherwise
b.grad.zero_()
print(w.item(), b.item()) # close to 3.0 and 5.0
Later, nn.Module and torch.optim will hide this boilerplate, but it is exactly what happens underneath.
Common mistakes
- Forgetting to zero gradients, so they accumulate across steps.
- Updating weights outside
torch.no_grad(), which makes autograd track the update itself. - Calling
.numpy()on a GPU tensor or on a tensor that requires grad (use.detach().cpu().numpy()). - Mixing
float64NumPy data withfloat32model weights — convert with.float().
Practice
Create a tensor t = torch.arange(24) and reshape it to (2, 3, 4). Compute the sum over the last dimension and the mean over the first. Then use autograd to find the derivative of f(x) = 5x^2 - 3x + 7 at x = 1.5 and confirm it by hand (answer: 12). Finally, modify the regression loop to use a learning rate of 0.1 and observe what happens to the loss.