PyTorch torch.optim Optimizer Module

Optimizers are core components in deep learning, responsible for adjusting model parameters based on the gradient of the loss function, allowing the model to gradually approach the optimal solution.

In PyTorch, the torch.optim module provides implementations of various optimization algorithms and is an indispensable tool for training neural networks.

Why Need Optimizers

Optimizers play a vital role in deep learning, solving the tedious problem of manually updating parameters.

  • Automated Parameter Updates: Manually calculating and updating each parameter is very tedious, and the optimizer does this automatically.
  • Accelerate Convergence: Using optimization algorithms finds the optimal solution faster than plain gradient descent.
  • Avoid Local Optima: Some optimizers have the ability to escape local optima.

Common Optimizer Types

Different optimizers are suitable for different scenarios. Choosing the right optimizer can significantly improve training results.

Optimizer NameMain FeaturesApplicable Scenarios
SGDSimple and basic, supports momentumBasic teaching, simple models, CNN
AdamAdaptive learning rateMost deep learning tasks
AdamWAdam + decoupled weight decayTasks requiring L2 regularization
RMSpropAdaptive learning rateRNN networks, speech recognition
AdagradPer-parameter learning rateSparse data, text processing
AdadeltaAdaptive learning rateLong-term training tasks

Core Optimizer API

Mastering the basic usage flow of optimizers is the first step in deep learning.

Basic Usage Flow

Optimizer usage follows a fixed pattern: create instance → zero gradients → backward pass → update parameters.

Example

import torch
import torch.nn as nn
import torch.optim as optim

# 1. Define a simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)

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

model = SimpleNet()

# 2. Create optimizer instance
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 3. Training loop
for epoch in range(epochs):
    # Forward pass
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    # Backward pass
    optimizer.zero_grad()  # Zero the gradient cache to avoid gradient accumulation
    loss.backward()        # Compute gradients

    # Parameter update
    optimizer.step()       # Update parameters

Key Methods Explanation

The optimizer provides several core methods to manage the parameter update process.

  • zero_grad(set_to_none=True): Zero the parameter gradient cache. When set to True, gradients are set to None, which saves more memory than setting them to 0.
  • step(): Perform a single parameter update, updating model parameters based on gradients and learning rate.
  • state_dict(): Get the optimizer state dictionary, which can be used to save checkpoints.
  • load_state_dict(state_dict): Load the optimizer state to resume training.
  • add_param_group(param_group): Dynamically add a parameter group.

Note: You must call zero_grad() before each backward pass, otherwise gradients will accumulate and cause unstable training. It is recommended to use zero_grad(set_to_none=True) to save memory.

Saving and Loading Optimizer State

When resuming training, both the model and optimizer states need to be saved.

Example

# Save checkpoint (save model, optimizer, and scheduler together)
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

# Load checkpoint
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1

Detailed Explanation of Common Optimizers

SGD (Stochastic Gradient Descent)

SGD is the most basic optimization algorithm, updating parameters by computing gradients from a single sample or a mini-batch of samples. It is the cornerstone of deep learning optimization, and many advanced optimizers are developed based on SGD.

Example

# SGD optimizer parameter description
# params: parameters to optimize (usually from model.parameters())
# lr: learning rate, controls the step size of parameter updates, default 0.01
# momentum: momentum factor, used to accelerate convergence and reduce oscillation, default 0
# weight_decay: L2 regularization coefficient, used to prevent overfitting, default 0
# dampening: momentum damping, controls the calculation of the momentum term, default 0
# nesterov: whether to use Nesterov momentum, default False
optimizer = optim.SGD(
    params=model.parameters(),
    lr=0.01,           # Learning rate
    momentum=0.9,      # Momentum factor
    weight_decay=1e-4, # L2 regularization
    nesterov=True      # Enable Nesterov momentum
)

Core parameter description:

  • lr (float): learning rate, controls the step size of parameter updates
  • momentum (float): momentum factor, used to accelerate convergence and reduce oscillation, common value 0.9
  • weight_decay (float): L2 regularization coefficient, used to prevent overfitting, common value 1e-4
  • nesterov (bool): whether to use Nesterov momentum, enabling it can reduce oscillation

Features:

  • Simple to implement, it is the basic algorithm of deep learning optimization
  • Adding a momentum term can accelerate convergence and improve training stability
  • Slower convergence speed, but the final accuracy may be higher
  • Suitable as a benchmark for comparison with other optimizers

Although SGD is simple, it can often achieve very good results with appropriate hyperparameters, making it a good starting point for learning optimization algorithms. In image classification tasks, SGD with momentum remains a mainstream choice.

Adam (Adaptive Moment Estimation)

Adam is one of the most commonly used optimizers, combining the advantages of momentum and adaptive learning rates. It adaptively adjusts the learning rate for each parameter by computing first-order and second-order moment estimates of the gradients.

Example

# Adam optimizer parameter description
# params: parameters to optimize
# lr: learning rate, default 0.001 (recommended value)
# betas: moving average coefficients used to compute gradients and squared gradients (beta1, beta2)
# beta1 controls the first-order moment estimate (momentum), default 0.9
# beta2 controls the second-order moment estimate (variance), default 0.999
# eps: numerical stability term, prevents division by zero, default 1e-8
# weight_decay: L2 regularization coefficient, default 0
# amsgrad: whether to use the AMSGrad variant, default False
optimizer = optim.Adam(
    params=model.parameters(),
    lr=0.001,                      # Recommended to use a smaller learning rate
    betas=(0.9, 0.999),            # Commonly used momentum parameters
    eps=1e-8,                      # Numerical stability term
    weight_decay=1e-4,             # L2 regularization
    amsgrad=False                  # Whether to use AMSGrad
)

Core parameter description:

  • betas (Tuple[float, float]): control the exponential moving average of gradients and squared gradients
  • eps (float): numerical stability term, prevents the denominator from being zero
  • amsgrad (bool): whether to use the AMSGrad variant, using it can guarantee convergence

Features:

  • Adaptive learning rate: automatically adjusts the learning rate based on the historical gradient of the parameter
  • Incorporates momentum concept: uses first-order moment estimate to accelerate convergence
  • Strong robustness: relatively insensitive to hyperparameter choices
  • Fast convergence speed, suitable for rapid prototyping

Adam is the default choice for most deep learning tasks, but in certain specific scenarios (such as GAN, reinforcement learning) you may need to try other optimizers.

AdamW（Adam with Weight Decay）

AdamW is an improved version of Adam that decouples weight decay from gradient updates, which is theoretically more conducive to convergence. In practical applications, AdamW usually performs better than Adam.

Example

# AdamW optimizer
# Main difference from Adam: different implementation of weight_decay
# AdamW's weight decay is more correct and does not affect gradient computation
optimizer = optim.AdamW(
    params=model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    weight_decay=0.01,    # Weight decay coefficient, usually set larger than with Adam
    amsgrad=False
)

# Recommended configuration: AdamW usually uses a weight_decay of 0.01
# While Adam usually uses 0.001

If your task requires weight decay (L2 regularization), it is strongly recommended to use AdamW instead of Adam.

RMSprop

RMSprop is an adaptive learning rate optimizer, especially suitable for handling non-stationary objectives and recurrent neural networks.

Example

# RMSprop optimizer
# Normalizes the learning rate by dividing by the exponentially weighted average of gradients
optimizer = optim.RMSprop(
    params=model.parameters(),
    lr=0.01,               # Learning rate
    alpha=0.99,            # Exponential decay rate for squared gradients
    eps=1e-8,              # Numerical stability term
    weight_decay=0,        # L2 regularization
    momentum=0,            # Momentum factor
    centered=False         # Whether to center the gradients
)

Adagrad

Adagrad is suitable for handling sparse data; it adaptively adjusts the learning rate for each parameter.

Example

# Adagrad optimizer
# Suitable for optimizing sparse data, uses a smaller learning rate for frequently updated parameters
optimizer = optim.Adagrad(
    params=model.parameters(),
    lr=0.01,               # Learning rate
    lr_decay=0,            # Learning rate decay
    weight_decay=0,       # L2 regularization
    initial_accumulator_value=0  # Initial accumulator value
)

Advanced Optimizer Techniques

Learning Rate Scheduling

Learning rate scheduling allows the learning rate to be dynamically adjusted during training, which can often significantly improve model convergence.

Example: Multiple Learning Rate Schedulers

from torch.optim.lr_scheduler import (
    StepLR,                # Step decay
    MultiStepLR,           # Multi-step decay
    ExponentialLR,         # Exponential decay
    CosineAnnealingLR,     # Cosine annealing
    ReduceLROnPlateau,     # Automatically adjust based on metrics
)

# Method 1: StepLR - decays every 30 epochs
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

# Method 2: MultiStepLR - decays at specified epochs
optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 80], gamma=0.1)

# Method 3: CosineAnnealingLR - cosine curve annealing
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)

# Method 4: ReduceLROnPlateau - automatically adjusts based on monitored metric
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = ReduceLROnPlateau(
    optimizer, mode='min',     # Monitor loss
    factor=0.5,                 # Decay factor
    patience=5,                 # Number of epochs to wait
    verbose=True                # Print information
)

# Training loop
for epoch in range(100):
    train_loss = train(...)
    val_loss = validate(...)

    # StepLR and other schedulers
    scheduler.step()

    # ReduceLROnPlateau requires the monitored metric to be passed in
    scheduler.step(val_loss)

The learning rate scheduler needs to be used together with the optimizer; step() must be called after optimizer.step(), otherwise the learning rate update may be abnormal.

Parameter Group Optimization

Parameter grouping allows different learning rates to be set for different layers, which is particularly useful in transfer learning.

Example

# Parameter grouping optimization example
# Set different learning rates for different layers
# Usually: the backbone network uses a smaller learning rate, and the classification head uses a larger learning rate
optimizer = optim.SGD([
    {'params': model.base.parameters(), 'lr': 1e-3},      # Base layers: large learning rate
    {'params': model.classifier.parameters(), 'lr': 1e-2} # Classification layer: large learning rate
], lr=1e-4)  # Global default learning rate (used when no parameter group is specified)

# More common notation in practical applications
optimizer = optim.Adam([
    {'params': model.fc.parameters(), 'lr': 1e-3},       # Classification head
    {'params': [p for n, p in model.named_parameters()    # Backbone network
                if not n.startswith('fc')],
     'lr': 1e-5},
])

Gradient Clipping

Gradient clipping can prevent gradient explosion and improve training stability. It is especially useful in deep networks such as RNN and LSTM.

Example

import torch.nn as nn

# Gradient clipping example
# max_norm: the maximum norm of gradients; gradients exceeding this value will be scaled
# norm_type: norm type, default is 2 (L2 norm)
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# Position in the training loop
for epoch in range(epochs):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()

    # Clip gradients after loss.backward() and before optimizer.step()
    nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    optimizer.step()

# Another way of gradient clipping: clipping by value
for param in model.parameters():
    if param.grad is not None:
        param.grad.data.clamp_(min=-1.0, max=1.0)

Gradient clipping is a common technique for training deep neural networks (especially RNN, LSTM, etc.), and can effectively prevent training crashes caused by gradient explosion.

Gradient Accumulation

When video memory is insufficient, gradient accumulation can be used to simulate the training effect of a large batch size.

Example

# Gradient accumulation example
# Actual batch_size = batch_size * accumulation_steps
accumulation_steps = 4   # Accumulate 4 small batches

optimizer.zero_grad()

for i, (inputs, labels) in enumerate(train_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)

    # Divide the loss by the accumulation steps to achieve averaging
    loss = loss / accumulation_steps
    loss.backward()

    # Update parameters once every specified number of accumulation steps
    if (i + 1) % accumulation_steps == 0:
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        optimizer.zero_grad()

# Handle the remaining gradients
if (i + 1) % accumulation_steps != 0:
    optimizer.step()
    optimizer.zero_grad()

Complete Training Example

The following is a complete training pipeline, demonstrating best practices for optimizers.

Example

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

# Configuration
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
EPOCHS = 100
BATCH_SIZE = 32
LR = 1e-3

# Create model and move to device
model = SimpleNet().to(DEVICE)

# Loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)

# Learning rate scheduler (cosine annealing)
scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS, eta_min=1e-6)

# Training loop
best_acc = 0.0
for epoch in range(EPOCHS):
    model.train()
    total_loss = 0.0
    correct = 0

    for inputs, labels in train_loader:
        inputs = inputs.to(DEVICE)
        labels = labels.to(DEVICE)

        # Clear gradients (recommended to use set_to_none=True)
        optimizer.zero_grad(set_to_none=True)

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

        # Backward pass
        loss.backward()

        # Gradient clipping (prevent gradient explosion)
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        # Update parameters
        optimizer.step()

        # Statistics
        total_loss += loss.item()
        correct += (outputs.argmax(1) == labels).sum().item()

    # Update learning rate
    scheduler.step()

    # Print training information
    avg_loss = total_loss / len(train_loader)
    accuracy = correct / len(train_loader.dataset)
    current_lr = scheduler.get_last_lr()[0]
    print(f"Epoch {epoch+1}/{EPOCHS} | Loss: {avg_loss:.4f} | "
          f"Acc: {accuracy:.4f} | LR: {current_lr:.6f}")

    # Save best model
    if accuracy > best_acc:
        best_acc = accuracy
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'best_acc': best_acc,
        }, 'best_model.pth')

print(f"Training complete, best accuracy: {best_acc:.4f}")

Optimizer Selection Guide

Choosing the right optimizer needs to be determined based on the specific task, data characteristics, and training stage.

Select by Task

Different task types have different optimizer recommendations.

Task TypeRecommended OptimizerRecommended Learning RateNotes
Image Classification (CNN)SGD + Momentum0.01 ~ 0.1Slow convergence but high accuracy
Image Classification (CNN)AdamW0.001Fast convergence
NLP / TransformerAdamW1e-5 ~ 1e-4Smaller learning rate
RNN / LSTMRMSprop / Adam0.001Adaptive learning rate
GANAdam (G) / Adam (D)0.0001Smaller learning rate
Reinforcement LearningAdam / RMSprop0.0001 ~ 0.001Depends on the specific task
Rapid experimentationAdam / AdamW0.001Fast convergence

Performance Comparison

OptimizerConvergence SpeedMemory UsageHyperparameter SensitivityFinal Accuracy
SGD + MomentumSlowLowHighHigh
AdamFastMediumLowMedium
AdamWFastMediumLowHigh
RMSpropMediumMediumMediumMedium
AdagradMediumHighMediumLow

Common Problems and Solutions

  • Unstable training (Loss oscillation): reduce learning rate, add gradient clipping
  • Convergence too slow: use Adam or AdamW, increase learning rate
  • Overfitting: increase weight_decay, use regularization
  • Insufficient video memory: reduce batch_size, use gradient accumulation

The choice of optimizer is not absolute; it is recommended to start with Adam or AdamW, and consider other optimizers if the results are not satisfactory. For specific tasks, experiments may be needed to determine the optimal choice.

For more detailed information, refer toPyTorch official documentation。

Other extensions