How to Use Adam Optimizer in PyTorch?

Deep Learning models use neural networks containing neurons within multiple layers. These neurons are connected across layers with their weights, and optimizers are responsible for updating those weights to reduce the loss value and improve accuracy. Two of the most popular optimizers in PyTorch are SGD and Adam.

What is the Adam Optimizer?

Adam (Adaptive Moment Estimation) controls the learning rate dynamically for each parameter in the model. Unlike SGD, which uses a fixed learning rate, Adam tracks the first moment (mean of gradients) and second moment (variance of gradients), and adjusts step sizes per parameter accordingly. This makes it converge faster and work well with minimal tuning for most deep learning tasks.

Syntax:

torch.optim.Adam(parameters, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

Adam Optimizer Parameters

Parameter	Default	Description
parameters	Required	Model parameters to optimize — use model.parameters()
lr	0.001	Learning rate — controls the step size at each iteration
betas	(0.9, 0.999)	Coefficients for computing running averages of gradient and its square
eps	1e-08	Small value for numerical stability — prevents division by zero
weight_decay	0	L2 regularization coefficient to reduce overfitting

How to Use Adam Optimizer in PyTorch

The following steps build a simple neural network and train it using the Adam optimizer.

Step 1: Import the Torch Library

import torch

Step 2: Set Up the Model’s Dimensions

batch = 128
input_dim = 2000
hidden_dim = 200
output_dim = 20

• batch — 128 mini-batches per epoch.
• input_dim — 2000 input features for the input layer.
• hidden_dim — 200 neurons in the hidden layer.
• output_dim — 20 outputs at the final layer.

Step 3: Build the Dataset

input = torch.randn(batch, input_dim)
output = torch.randn(batch, output_dim)

Step 4: Build the Neural Network Model

model = torch.nn.Sequential(
    torch.nn.Linear(input_dim, hidden_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(hidden_dim, output_dim),
)

The Sequential model passes data through layers in order: a Linear input layer, a ReLU activation, and a final Linear output layer.

Step 5: Initialize the Loss Function and Adam Optimizer

loss_fn = torch.nn.MSELoss(reduction='sum')
optim = torch.optim.Adam(model.parameters(), lr=0.05)

Step 6: Train the Model

for epoch in range(1):

    running_loss = 0
    for i, data in enumerate(input, 0):

        optim.zero_grad()
        pred = model(input)
        loss = loss_fn(pred, output)
        loss.backward()
        optim.step()
        running_loss += loss.item()
        print('[%d, %5d] Value: %.3f' %
            (epoch + 1, i + 1, running_loss / 10000))
        running_loss = 0

print('Finished Training')

• zero_grad() — clears gradients accumulated from the previous step.
• loss.backward() — computes gradients via backpropagation.
• optim.step() — updates the model’s weights using Adam’s adaptive rule.

The decreasing loss values confirm the Adam optimizer is improving the model with each batch.

Adam vs SGD: When to Use Which

Feature	Adam	SGD
Learning rate	Adaptive per parameter	Fixed (same for all parameters)
Convergence speed	Fast — good default for most tasks	Slower but can generalize better
Memory usage	Higher (stores gradient moments)	Lower
Momentum	Built-in adaptive momentum	Optional manual momentum parameter
Best for	RNNs, Transformers, NLP, most DL tasks	CNNs, image classification, fine-tuning

Adam is the recommended default for most deep learning projects. Consider switching to SGD with momentum when you need stronger generalization, particularly for image classification tasks like CIFAR-10 where SGD has historically achieved better final accuracy.

Conclusion

To use the Adam optimizer in PyTorch, initialize it with torch.optim.Adam(model.parameters(), lr=0.001) and call zero_grad(), loss.backward(), and optimizer.step() in each training iteration. Adam’s adaptive learning rate makes it the go-to optimizer for most deep learning workflows — it converges faster than SGD and requires minimal hyperparameter tuning. Adjust the lr, betas, and weight_decay parameters based on the specific requirements of your model.