pytorch-lr-finder

When using LRFinder, keep the following constraints and behaviors in mind:

• Optimizer State: The optimizer passed to LRFinder must not have an LRScheduler attached to it.
• Weight Modification: LRFinder.range_test() modifies model weights and optimizer parameters. Always call lr_finder.reset() to restore them to their initial state.
• Data Format: range_test() expects DataLoader objects to return a pair of (input, label). The input must be ready for the model and label must be ready for the criterion without further processing. If processing is required, use the TrainDataLoaderIter and ValDataLoaderIter wrapper classes.
• History Access: You can access the results via lr_finder.history, which returns a dictionary with lr and loss keys.

To run the provided examples, you must have Jupyter Notebook installed along with torchvision and ipywidgets.

pip install torchvision
pip install ipywidgets

If your hardware constraints require a smaller batch size than your desired training batch size, use the accumulation_steps parameter in LRFinder.range_test().

Note: Ensure the batch_size in your DataLoader is set to the smaller real_batch_size used for the actual iterations.

from torch.utils.data import DataLoader
from torch_lr_finder import LRFinder

desired_batch_size, real_batch_size = 32, 4
accumulation_steps = desired_batch_size // real_batch_size

dataset = ...
trainloader = DataLoader(dataset, batch_size=real_batch_size, shuffle=True)

model = ...
criterion = ...
optimizer = ...

lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=10, num_iter=100, step_mode="exp", accumulation_steps=accumulation_steps)
lr_finder.plot()
lr_finder.reset()

Install the package using pip. To include support for mixed precision training (Apex), use the --global-option="apex" flag during installation.

# Standard installation
pip install torch-lr-finder

# Installation with mixed precision (Apex) support
pip install torch-lr-finder -v --global-option="apex"

If the torch_lr_finder package is not installed in your environment, you can import it directly from the source by adding the repository root to your sys.path.

try:
    from torch_lr_finder import LRFinder
except ImportError:
    import sys
    sys.path.insert(0, '..')
    from torch_lr_finder import LRFinder

To use PyTorch's native torch.amp, provide an amp_config dictionary (containing device_type and dtype) and a grad_scaler instance to the LRFinder constructor, and set amp_backend='torch'.

from torch_lr_finder import LRFinder

amp_config = {
    'device_type': 'cuda',
    'dtype': torch.float16,
}
grad_scaler = torch.cuda.amp.GradScaler()

lr_finder = LRFinder(
    model, optimizer, criterion, device='cuda',
    amp_backend='torch', amp_config=amp_config, grad_scaler=grad_scaler
)
lr_finder.range_test(trainloader, end_lr=10, num_iter=100, step_mode='exp')
lr_finder.plot()
lr_finder.reset()

To use NVIDIA's apex.amp, initialize your model and optimizer with amp.initialize before passing them to LRFinder, and set amp_backend='apex'.

from torch_lr_finder import LRFinder
from apex import amp

# Add this line before running `LRFinder`
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')

lr_finder = LRFinder(model, optimizer, criterion, device='cuda', amp_backend='apex')
lr_finder.range_test(trainloader, end_lr=10, num_iter=100, step_mode='exp')
lr_finder.plot()
lr_finder.reset()

This approach increases the learning rate linearly and computes evaluation loss. It typically produces more precise curves because evaluation loss is more sensitive to divergence, but it is slower. When using step_mode="linear", ensure the learning rate range is within the same order of magnitude.

from torch_lr_finder import LRFinder

model = ...
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.1, weight_decay=1e-2)

lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, val_loader=val_loader, end_lr=1, num_iter=100, step_mode="linear")
lr_finder.plot(log_lr=False)
lr_finder.reset()

This approach increases the learning rate exponentially and computes training loss. It is the default behavior. Use lr_finder.plot() to visualize the training loss versus the logarithmic learning rate.

from torch_lr_finder import LRFinder

model = ...
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-7, weight_decay=1e-2)

lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=100, num_iter=100)
lr_finder.plot() # to inspect the loss-learning rate graph
lr_finder.reset() # to reset the model and optimizer to their initial state

To perform a learning rate range test using the training loss (following the fastai procedure), initialize LRFinder with your model, optimizer, and criterion. Call range_test() on your training dataloader.

Key parameters:

• end_lr: The maximum learning rate to reach.
• num_iter: The number of iterations to run the test.
• step_mode: Use "linear" for small ranges or "exp" (exponential) for larger ranges.

After the test, use .plot() to visualize the training loss vs. learning rate. Use .reset() to restore the model and optimizer to their initial states.

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-7, weight_decay=1e-2)
lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=100, num_iter=100, step_mode="exp")
lr_finder.plot()
lr_finder.reset()