Interactive PyTorch Cheatsheet

An interactive guide to PyTorch concepts, neural network building, and deep learning patterns.

0. Core PyTorch Concepts

Tensors

  • The fundamental data structure in PyTorch. Similar to NumPy arrays but can run on GPUs.
  • When to Use: Tensors are used for all data representation in PyTorch, from input data and model parameters to intermediate computations and outputs.
  • Creation:
    • torch.tensor([1.0, 2.0])
    • torch.zeros(3, 3)
    • torch.ones(2, 4)
    • torch.rand(2, 2)
  • Attributes: x.shape, x.dtype, x.device
  • Operations: Element-wise (+, -, *, /), Matrix multiplication (x.matmul(y)), Transpose (x.T), Reshape (x.view(-1)), Aggregations (x.sum(), x.mean(), x.max())

Autograd (Automatic Differentiation)

  • PyTorch's engine for automatic computation of gradients. Essential for backpropagation.
  • When to Use: Automatically tracks operations on tensors to compute gradients during the backward pass, which are then used by optimizers to update model parameters.
  • Enable gradient tracking: x = torch.tensor([2.0], requires_grad=True)
  • Compute gradients: loss.backward()
  • Access gradients: x.grad
  • Disable gradient tracking for inference: with torch.no_grad():

GPU Support

  • When to Use: When your machine has an NVIDIA GPU and you want to significantly speed up computations for large models and datasets.
  • Check availability: torch.cuda.is_available()
  • Set device: device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  • Move tensor/model to device: tensor.to(device), model.to(device)

1. Building Neural Networks (torch.nn)

Neural networks in PyTorch are typically built by subclassing torch.nn.Module.

Basic Structure

import torch
import torch.nn as nn
import torch.nn.functional as F

class MyNeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MyNeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size) # Fully connected layer
        self.relu = nn.ReLU() # Activation function
        self.fc2 = nn.Linear(hidden_size, output_size)

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

# Or using nn.Sequential for simple linear stacks:
model_sequential = nn.Sequential(
    nn.Linear(input_size, hidden_size),
    nn.ReLU(),
    nn.Linear(hidden_size, output_size)
)

Types of Layers (torch.nn)

Linear/Dense Layers (nn.Linear)

  • When to Use: For applying a linear transformation to input data. These are the workhorses of fully connected neural networks and are often used as output layers or in combination with other layers. Ideal for tabular data or the final layers of image/sequence models after feature extraction.
  • nn.Linear(in_features, out_features, bias=True): Applies a linear transformation (\(y = xA^T + b\)).

Convolutional Layers (nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose)

  • When to Use: Primarily for processing data with a known grid-like topology, such as images (2D conv) or time series/sequences (1D conv). They excel at learning local patterns and hierarchies of features. ConvTranspose (deconvolution) is used for upsampling, e.g., in generative models or image segmentation.
  • nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0): 1D convolution.
  • nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0): 2D convolution.
  • nn.Conv3d(...): 3D convolution.
  • nn.ConvTranspose1d/2d/3d(...): Transposed convolutions (for upsampling).

Pooling Layers (nn.MaxPool, nn.AvgPool)

  • When to Use: Typically used after convolutional layers to reduce the spatial dimensions (height and width for images), which helps to reduce computational cost, control overfitting, and make the model more robust to small variations in input.
  • nn.MaxPool1d/2d/3d(kernel_size, stride=None, padding=0): Max pooling.
  • nn.AvgPool1d/2d/3d(kernel_size, stride=None, padding=0): Average pooling.

Recurrent Layers (nn.RNN, nn.LSTM, nn.GRU)

  • When to Use: For processing sequential data where the order of elements matters, such as natural language, time series, or audio. LSTMs and GRUs are more common than basic RNNs as they address the vanishing/exploding gradient problem and better capture long-range dependencies.
  • nn.RNN(input_size, hidden_size, num_layers=1, ...): Simple RNN.
  • nn.LSTM(input_size, hidden_size, num_layers=1, ...): Long Short-Term Memory.
  • nn.GRU(input_size, hidden_size, num_layers=1, ...): Gated Recurrent Unit.

Activation Functions (nn.ReLU, nn.Sigmoid, nn.Tanh, nn.Softmax, etc.)

  • When to Use: Applied after linear transformations within layers to introduce non-linearity, allowing the network to learn complex patterns and relationships that linear models cannot.
  • nn.ReLU(): Most common choice for hidden layers due to computational efficiency and mitigating vanishing gradients.
  • nn.LeakyReLU(), nn.ELU(): Alternatives to ReLU to prevent "dying ReLU" problem.
  • nn.Sigmoid(): Useful for binary classification output layers (output between 0 and 1) or in gates of LSTMs/GRUs.
  • nn.Tanh(): Output between -1 and 1, can be used in hidden layers or recurrent networks.
  • nn.Softmax(dim): Used in the output layer for multi-class classification to convert raw scores (logits) into probabilities that sum to 1. Note: nn.CrossEntropyLoss already includes Softmax internally.

Normalization Layers (nn.BatchNorm, nn.LayerNorm, nn.GroupNorm)

  • When to Use: To stabilize and accelerate training by normalizing the activations of previous layers. This helps in dealing with internal covariate shift.
  • nn.BatchNorm1d/2d/3d(): Most common for normalizing across the batch dimension, typically used in feedforward or convolutional networks.
  • nn.LayerNorm(): Normalizes across the feature dimension, often preferred in recurrent networks and Transformers.
  • nn.GroupNorm(): A flexible alternative that normalizes features within groups, useful when batch size is very small.

Dropout Layers (nn.Dropout)

  • When to Use: A regularization technique to prevent overfitting. Randomly sets a fraction (p) of input units to zero during training, forcing the network to learn more robust features.
  • nn.Dropout(p=0.5)

Flatten Layer (nn.Flatten)

  • When to Use: To convert multi-dimensional feature maps (e.g., from convolutional layers) into a 1D vector before feeding them into a linear layer.
  • nn.Flatten(start_dim=1)

2. Loss Functions (torch.nn)

Quantify the difference between predicted and actual values. Choosing the right loss function is crucial for effective training.

Regression Losses

  • nn.MSELoss() (Mean Squared Error / L2 Loss):
    • When to Use: For regression tasks where the goal is to predict a continuous numerical value. It heavily penalizes larger errors.
    • Formula: \( \frac{1}{N} \sum_{i=1}^N (y_i - \hat{y}_i)^2 \)
  • nn.L1Loss() (Mean Absolute Error / L1 Loss):
    • When to Use: Also for regression tasks. It is less sensitive to outliers than MSELoss because it takes the absolute difference instead of the squared difference.
    • Formula: \( \frac{1}{N} \sum_{i=1}^N |y_i - \hat{y}_i| \)
  • nn.SmoothL1Loss() (Huber Loss):
    • When to Use: A hybrid of L1 and L2 loss. It behaves like MSE for small errors and L1 for large errors, making it robust to outliers while still providing a smooth gradient around the optimum. Useful when some outliers are expected.

Classification Losses

  • nn.CrossEntropyLoss():
    • When to Use: The standard for multi-class classification problems where the target is a single class index (e.g., predicting digit 0-9). It expects raw, unnormalized scores (logits) from the network's output layer. It combines LogSoftmax and NLLLoss for numerical stability.
  • nn.NLLLoss() (Negative Log Likelihood Loss):
    • When to Use: Used for multi-class classification if your model's output layer explicitly applies nn.LogSoftmax. If your model outputs logits directly, nn.CrossEntropyLoss is preferred.
  • nn.BCELoss() (Binary Cross Entropy Loss):
    • When to Use: For binary classification (two classes, 0 or 1) when the model's output is a probability (between 0 and 1) obtained after a Sigmoid activation function.
  • nn.BCEWithLogitsLoss():
    • When to Use: The preferred choice for binary classification. It is numerically more stable than applying Sigmoid and then BCELoss separately, as it computes the sigmoid within the loss function itself using a log-sum-exp trick. It expects raw logits as input.
  • nn.MultiLabelSoftMarginLoss() or nn.BCEWithLogitsLoss() (with multiple outputs):
    • When to Use: For multi-label classification problems where an instance can belong to multiple classes simultaneously (e.g., an image can contain both "dog" and "cat"). For BCEWithLogitsLoss, you'd typically have one output neuron per class and apply it across all of them.

3. Optimizers (torch.optim)

Algorithms to update model weights and biases based on gradients to minimize the loss.

  • Initialization: optimizer = optim.OptimizerName(model.parameters(), lr=learning_rate, ...)

Common Optimizers

  • optim.SGD(params, lr, momentum=0, weight_decay=0): Stochastic Gradient Descent
    • When to Use: A foundational optimizer. While simple, with carefully tuned learning rates and momentum, it can achieve excellent results and often generalizes well. It's a good baseline to start with. Momentum helps accelerate SGD in the relevant direction and dampens oscillations.
  • optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0): Adaptive Moment Estimation
    • When to Use: One of the most popular and generally effective optimizers for a wide range of deep learning tasks. It adapts the learning rate for each parameter individually and combines concepts of momentum and RMSprop. Often a good starting point if you're unsure which optimizer to pick.
  • optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0): Root Mean Square Propagation.
    • When to Use: Useful for recurrent neural networks and when dealing with non-stationary objectives. It divides the learning rate by an exponentially decaying average of squared gradients.
  • optim.Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0, initial_accumulator_value=0, eps=1e-10): Adaptive Gradient Algorithm.
    • When to Use: Suitable for sparse data, as it adapts the learning rate to features, giving larger updates to infrequent features. However, its learning rate can decay too aggressively, leading to slow training later on.
  • optim.Adadelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0): Adadelta.
    • When to Use: An extension of Adagrad that aims to overcome its aggressively diminishing learning rates by limiting the window of accumulated past gradients. It doesn't require a learning rate to be set manually.

Learning Rate Schedulers (torch.optim.lr_scheduler)

  • When to Use: To dynamically adjust the learning rate during training. This can significantly improve model performance and training stability.
  • torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1):
    • When to Use: For a fixed, stepwise decay of the learning rate. For example, decay LR by a factor of 0.1 every step_size epochs. Simple and effective for many tasks.
  • torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=10):
    • When to Use: When you want to reduce the learning rate adaptively. It monitors a quantity (e.g., validation loss) and reduces the learning rate when that quantity has stopped improving for a patience number of epochs. Very common for achieving better convergence.
  • Other schedulers exist (e.g., ExponentialLR, CosineAnnealingLR) for more specific decay patterns.

4. Data Handling

torch.utils.data.Dataset

  • When to Use: When you have custom data (not covered by torchvision or torchtext datasets) and need to define how samples are loaded and preprocessed. You subclass Dataset and implement __len__ (returns dataset size) and __getitem__ (loads and returns a sample by index).

torch.utils.data.DataLoader

  • When to Use: To efficiently load data in batches, shuffle it, and optionally use multiple processes (num_workers) for parallel data loading. It abstracts away batching, shuffling, and multi-threading.
  • Key arguments: dataset, batch_size, shuffle=True, num_workers (for parallel data loading).

torchvision.datasets and torchvision.transforms

  • When to Use: For common computer vision tasks. torchvision.datasets provides ready-to-use datasets (e.g., MNIST, CIFAR10). torchvision.transforms offers a wide array of image transformations (resizing, cropping, normalization, data augmentation) to prepare images for neural networks.
  • transforms.ToTensor(): Converts PIL Image or NumPy array to torch.FloatTensor.
  • transforms.Normalize((mean,), (std,)): Normalizes a tensor image with mean and standard deviation.

5. Training Loop Essentials

# 1. Instantiate Model, Loss Function, and Optimizer
model = MyNeuralNetwork(input_size, hidden_size, output_size).to(device)
criterion = nn.CrossEntropyLoss() # Or other appropriate loss
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 2. Training Loop
num_epochs = 10
for epoch in range(num_epochs):
    model.train() # Set model to training mode
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device), labels.to(device)

        # Zero the parameter gradients
        optimizer.zero_grad()

        # Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # Backward pass and optimize
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {running_loss / len(train_loader):.4f}")

    # 3. Evaluation (Optional, usually after each epoch or few epochs)
    model.eval() # Set model to evaluation mode
    correct = 0
    total = 0
    with torch.no_grad(): # Disable gradient calculations for inference
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1) # Get the predicted class
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f"Accuracy on test set: {accuracy:.2f}%")

6. Saving and Loading Models

  • When to Use: To persist trained models for later inference, continued training, or deployment.
  • Saving the entire model: (Not recommended for production as it relies on the exact class definition being available at load time)
    • torch.save(model, 'model.pth')
  • Loading the entire model:
    • model = torch.load('model.pth')
  • Saving only the model's state dictionary (recommended):
    • When to Use: This is the standard and recommended way. It saves only the learned parameters (weights and biases), making the saved file smaller and more flexible for loading into different model architectures (as long as layer names match) or different PyTorch versions.
    • torch.save(model.state_dict(), 'model_state_dict.pth')
  • Loading the model's state dictionary:
    model = MyNeuralNetwork(input_size, hidden_size, output_size) # Re-instantiate the model with its class definition
    model.load_state_dict(torch.load('model_state_dict.pth'))
    model.eval() # Set to evaluation mode after loading for inference

7. Best Practices

  • model.train() vs. model.eval():
    • When to Use: Essential for correctly handling layers like Dropout and BatchNorm.
      • model.train(): Sets the module in training mode (e.g., enables dropout and batch normalization updates to their running statistics).
      • model.eval(): Sets the module in evaluation mode (e.g., disables dropout, uses fixed running means/variances for batch normalization, no updates to statistics).
  • optimizer.zero_grad():
    • When to Use: At the beginning of each training iteration (batch). Gradients accumulate by default in PyTorch, so you must zero them out before computing new gradients for the current batch.
  • with torch.no_grad():
    • When to Use: During validation, testing, and inference. This disables gradient calculations, which saves memory and speeds up computations because PyTorch doesn't need to build the computational graph for backpropagation.
  • Reproducibility:
    • When to Use: When you need your training runs to be repeatable, for debugging, research, or comparing models.
    • torch.manual_seed(seed)
    • np.random.seed(seed)
    • torch.backends.cudnn.deterministic = True
    • torch.backends.cudnn.benchmark = False
  • Early Stopping:
    • When to Use: To prevent overfitting and save training time. Monitor a metric (e.g., validation loss or accuracy) and stop training if it doesn't improve for a certain number of epochs.
  • Learning Rate Scheduling:
    • When to Use: Almost always. A fixed learning rate is rarely optimal throughout the entire training process. Schedulers allow for dynamic adjustment, often leading to faster convergence and better final performance.
  • Weight Decay (L2 Regularization):
    • When to Use: To prevent overfitting. It adds a penalty to the loss function that is proportional to the square of the magnitude of the weights. Most optimizers (like SGD, Adam) have a weight_decay parameter.
  • Data Augmentation:
    • When to Use: Especially for image (or audio/text) tasks when your dataset size is limited. It artificially increases the diversity of your training data by applying random transformations, making the model more robust.
  • Debugging:
    • When to Use: Throughout the model development and training process.
    • Print tensor shapes (.shape) and values: Crucial for understanding data flow and identifying dimension mismatches.
    • Check requires_grad attribute of tensors: To ensure gradients are being tracked where needed.
    • Use torchsummary (external library): from torchsummary import summary; summary(model, input_size=(channels, height, width)) - provides a quick overview of model layers, output shapes, and parameter counts.
  • Gradient Clipping:
    • When to Use: For recurrent neural networks (RNNs, LSTMs, GRUs) or very deep networks where exploding gradients can occur. It caps the magnitude of gradients during backpropagation to prevent them from becoming too large and causing unstable training.
    • nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    • nn.utils.clip_grad_value_(model.parameters(), clip_value)