PyTorch 2 cheatsheet

PyTorch 2: Autograd and models

5 common PyTorch 2 tasks.

Compute gradients

Track operations and compute derivatives with backpropagation.

w = torch.tensor(2.0, requires_grad=True)
loss = (w * 3 - 1) ** 2
loss.backward()
print(w.grad)

Skip gradient tracking

Save memory and time when predicting.

with torch.no_grad():
    preds = model(x)

Define a model

Subclass nn.Module, create layers in __init__ and define forward.

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.layers(x)

Count a model's parameters

Count trainable parameters.

sum(p.numel() for p in model.parameters() if p.requires_grad)

Load data in batches

Wrap a dataset to get shuffled mini-batches.

from torch.utils.data import DataLoader

loader = DataLoader(dataset, batch_size=64, shuffle=True)