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.CrossEntropyLossalready 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
LogSoftmaxandNLLLossfor numerical stability.
- 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
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.CrossEntropyLossis preferred.
- When to Use: Used for multi-class classification if your model's output layer explicitly applies
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
Sigmoidactivation function.
- 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
nn.BCEWithLogitsLoss():- When to Use: The preferred choice for binary classification. It is numerically more stable than applying
Sigmoidand thenBCELossseparately, as it computes the sigmoid within the loss function itself using a log-sum-exp trick. It expects raw logits as input.
- When to Use: The preferred choice for binary classification. It is numerically more stable than applying
nn.MultiLabelSoftMarginLoss()ornn.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.
- 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
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_sizeepochs. Simple and effective for many tasks.
- When to Use: For a fixed, stepwise decay of the learning rate. For example, decay LR by a factor of 0.1 every
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
patiencenumber of epochs. Very common for achieving better convergence.
- 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
- 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
torchvisionortorchtextdatasets) and need to define how samples are loaded and preprocessed. You subclassDatasetand 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.datasetsprovides ready-to-use datasets (e.g., MNIST, CIFAR10).torchvision.transformsoffers 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 totorch.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
DropoutandBatchNorm.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).
- When to Use: Essential for correctly handling layers like
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 = Truetorch.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_decayparameter.
- 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
- 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_gradattribute 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)