← All cheatsheets

PyTorch Cheatsheet

Tensors, autograd, nn.Module, loss functions, training loops, and classification metrics.

Tensors — Creation & Attributes

import torch

torch.tensor([1, 2, 3])                 # from a Python list
torch.zeros(3); torch.ones(3); torch.full((3,), 7)
torch.arange(0, 10, 2); torch.linspace(0, 1, 5)
torch.rand(3, 3)                        # uniform [0, 1)
torch.randn(3, 3)                       # standard normal
torch.zeros_like(x); torch.ones_like(x)

x.shape; x.dtype; x.device; x.ndim; x.numel()
x.float(); x.long(); x.to(torch.float32)   # dtype casts
x.to(device); x.cuda(); x.cpu()

Shape Manipulation

x.reshape(2, 3); x.view(2, 3)           # view requires contiguous memory
x.squeeze()                              # remove all size-1 dims
x.unsqueeze(dim)                         # insert a new size-1 dim at position dim
x.permute(2, 0, 1)                       # reorder dims arbitrarily (e.g. HWC -> CHW)
x.t(); x.transpose(0, 1)                 # swap two dims
torch.cat([a, b], dim=0)                 # join along an EXISTING dimension
torch.stack([a, b], dim=0)               # join along a NEW dimension
torch.flatten(x)                         # collapse to 1-D
torch.clamp(x, lo, hi)                   # clip values to a range

Indexing & Math

x[0]; x[:, 0]; x[x > 0]                  # same semantics as NumPy
x.argmin(); x.argmax()                   # positions of min/max
x.min(); x.max(); x.sum(); x.mean(); x.std()
x @ y; torch.matmul(x, y)                # matrix multiply
x * y                                     # elementwise (broadcasting rules match NumPy)
torch.softmax(x, dim=1)                  # softmax along a given dimension

Autograd

w = torch.tensor(1.0, requires_grad=True)   # track gradients for this tensor
loss = ((w * x - y) ** 2).mean()
loss.backward()                             # populate .grad on all leaf tensors
w.grad                                      # the computed gradient
w.grad.zero_()                              # reset before the next backward pass
with torch.no_grad():                       # disable tracking (e.g. during eval/updates)
    w -= lr * w.grad
x.detach()                                  # a tensor sharing data but detached from the graph

Only leaf tensors with requires_grad=True accumulate .grad; gradients accumulate across calls unless you zero them (optimizer.zero_grad() in a loop).

Building Models (torch.nn)

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(in_features, out_features)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(out_features, num_classes)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = Net()
model(x)                    # forward pass — same as model.forward(x)

nn.Linear(in, out) computes y = x @ W.T + b internally. Common layers: nn.Conv2d, nn.ReLU, nn.Sigmoid, nn.Dropout, nn.BatchNorm1d/2d, nn.Sequential(layer1, layer2, ...) for a simple stack.

Loss Functions

nn.MSELoss()(pred, target)              # mean squared error (regression)
nn.L1Loss()(pred, target)               # mean absolute error
nn.BCEWithLogitsLoss()(logits, target)  # binary classification, from raw logits
nn.CrossEntropyLoss()(logits, target)   # multi-class, target = integer class index

Training Loop Skeleton

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)   # or torch.optim.Adam
loss_fn = nn.CrossEntropyLoss()

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    preds = model(X_train)
    loss = loss_fn(preds, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        test_preds = model(X_test)
        test_loss = loss_fn(test_preds, y_test)

Classification Metrics (from raw predictions)

preds = torch.round(torch.sigmoid(logits))        # binary: logits -> 0/1 labels
preds = torch.softmax(logits, dim=1).argmax(dim=1) # multi-class: logits -> class index
accuracy = (preds == y_true).float().mean() * 100
tp = ((preds == 1) & (y_true == 1)).sum().float()
fp = ((preds == 1) & (y_true == 0)).sum().float()
fn = ((preds == 0) & (y_true == 1)).sum().float()
precision = tp / (tp + fp); recall = tp / (tp + fn)
f1 = 2 * precision * recall / (precision + recall)

Saving & Loading

torch.save(model.state_dict(), "model.pt")
model.load_state_dict(torch.load("model.pt"))

Gotchas

Also see: NumPy, Pandas