PyTorch Optimizer Collection

0.3.0 · active · verified Thu Apr 16

pytorch-optimizer is a production-focused optimization toolkit for PyTorch, offering a comprehensive collection of over 100 optimizers, 10+ learning rate schedulers, and 10+ loss functions. It provides a consistent API for fast experimentation with modern training methods without extensive boilerplate. The library is currently at version 0.3.0 and has seen periodic updates, with the latest release in October 2021.

Common errors

```
ImportError: cannot import name 'RAdam' from 'torch_optimizer'
```
cause Attempting to import RAdam while `pytorch-optimizer==0.2.0` is installed, where RAdam was temporarily removed.

fix Upgrade the package to `pip install pytorch-optimizer>=0.3.0` or downgrade to `pip install pytorch-optimizer==0.1.0`. Alternatively, use `torch.optim.RAdam` if your PyTorch version includes it.
```
ModuleNotFoundError: No module named 'pytorch_optimizer'
```
cause Incorrect module name used in the import statement. The PyPI package name (`pytorch-optimizer`) differs from the Python import module name (`torch_optimizer`).

fix Change the import statement to `import torch_optimizer as optim` or `from torch_optimizer import ...`.
```
TypeError: optimizer() missing 1 required positional argument: 'params'
```
cause When initializing an optimizer directly (e.g., `AdamP(...)`), the `model.parameters()` iterable of parameters to optimize was not passed.

fix Ensure that `model.parameters()` is passed as the first argument when instantiating any optimizer: `optimizer = AdamP(model.parameters(), lr=0.001)`.

Warnings

breaking The RAdam optimizer was temporarily removed in version `0.2.0` due to its inclusion in PyTorch core, causing `ImportError` for existing users. It was subsequently re-added in version `0.3.0` ('Revert for Drop RAdam'). Users on `0.2.0` will not have RAdam available directly from `torch_optimizer`.
Fix: Upgrade to `pytorch-optimizer>=0.3.0` to regain RAdam, or switch to `torch.optim.RAdam` for versions where it's available in PyTorch.
gotcha The PyPI package is named `pytorch-optimizer`, but the correct Python module to import is `torch_optimizer`. Attempting to `import pytorch_optimizer` will result in an `ModuleNotFoundError`.
Fix: Always use `import torch_optimizer` or `from torch_optimizer import ...` for importing classes and functions.
gotcha The library author advises against selecting optimizers solely based on visualizations. Different optimizers have unique properties and may require specific learning rate schedules or tuning. It is recommended to start with built-in PyTorch optimizers like SGD or Adam to establish a baseline before experimenting with `pytorch-optimizer` variants.

Install

pip install pytorch-optimizer Install stable version

Imports

AdamP
wrong:
```
from pytorch_optimizer import AdamP
```
correct:
```
from torch_optimizer import AdamP
```
The PyPI package name is `pytorch-optimizer`, but the import module is `torch_optimizer`.
load_optimizer
wrong:
```
from pytorch_optimizer import load_optimizer
```
correct:
```
from torch_optimizer import load_optimizer
```
The PyPI package name is `pytorch-optimizer`, but the import module is `torch_optimizer`.
create_optimizer
wrong:
```
from pytorch_optimizer import create_optimizer
```
correct:
```
from torch_optimizer import create_optimizer
```
The PyPI package name is `pytorch-optimizer`, but the import module is `torch_optimizer`.

Quickstart

This quickstart demonstrates how to initialize a model and apply an optimizer from the `torch_optimizer` library, using AdamP as an example. It includes a basic forward and backward pass to illustrate its use in a typical PyTorch training loop. You can directly import specific optimizers or use `load_optimizer` and `create_optimizer` for dynamic loading.

import torch
import torch.nn as nn
from torch_optimizer import AdamP

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

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

model = SimpleModel()

# Initialize AdamP optimizer with model parameters
# Replace with other optimizers like A2GradExp, AdaBelief, etc.
optimizer = AdamP(model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=1e-2)

# Example of a training step
inputs = torch.randn(32, 10) # Batch of 32, 10 features
targets = torch.randn(32, 1) # Corresponding targets

optimizer.zero_grad()
outputs = model(inputs)
loss = torch.nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

print(f"Initial loss: {loss.item():.4f}")

view raw JSON →