How to use FastaiLRFinder with Ignite
This how-to guide demonstrates how we can leverage the
FastaiLRFinder handler to find an optimal learning rate to train our model on. We will compare the results produced with and without using the handler for better understanding.
In this example, we will be using a ResNet18 model on the MNIST dataset. The base code is the same as used in the Getting Started Guide.
Basic Setup
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.models import resnet18
from torchvision.transforms import Compose, Normalize, ToTensor
from ignite.engine import create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Accuracy, Loss
from ignite.handlers import FastaiLRFinder
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.model = resnet18(num_classes=10)
self.model.conv1 = nn.Conv2d(
1, 64, kernel_size=3, padding=1, bias=False
)
def forward(self, x):
return self.model(x)
model = Net().to(device)
data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
train_loader = DataLoader(
MNIST(download=True, root=".", transform=data_transform, train=True),
batch_size=128,
shuffle=True,
)
test_loader = DataLoader(
MNIST(download=True, root=".", transform=data_transform, train=False),
batch_size=256,
shuffle=False,
)
model = Net().to(device)
optimizer = torch.optim.RMSprop(model.parameters(), lr=1e-06)
criterion = nn.CrossEntropyLoss()
We will first train the model with a fixed learning rate (lr) of 1e-06 and inspect our results. Let’s save the initial state of the model and the optimizer to restore them later for comparison.
init_model_state = model.state_dict()
init_opt_state = optimizer.state_dict()
Without LR Finder
trainer = create_supervised_trainer(model, optimizer, criterion, device=device)
trainer.run(train_loader, max_epochs=3)
State:
iteration: 1407
epoch: 3
epoch_length: 469
max_epochs: 3
output: 0.5554001927375793
batch: <class 'list'>
metrics: <class 'dict'>
dataloader: <class 'torch.utils.data.dataloader.DataLoader'>
seed: <class 'NoneType'>
times: <class 'dict'>
evaluator = create_supervised_evaluator(
model, metrics={"Accuracy": Accuracy(), "Loss": Loss(criterion)}, device=device
)
evaluator.run(test_loader)
print(evaluator.state.metrics)
{'Accuracy': 0.8655, 'Loss': 0.602867822265625}
Let’s see how we can achieve better results by using the
FastaiLRFinder handler. But first, let’s restore the initial state of the model and optimizer so we can re-train them from scratch.
model.load_state_dict(init_model_state)
optimizer.load_state_dict(init_opt_state)
With LR Finder
When attached to the trainer, this handler follows the same procedure used by
fastai. The model is trained for num_iter iterations while the learning rate is increased from start_lr (defaults to initial value specified by the optimizer, here 1e-06) to the upper bound called end_lr. This increase can be linear (step_mode="linear") or exponential (step_mode="exp"). The default step_mode is exponential which is recommended for larger learning rate ranges while linear provides good results for small ranges.
lr_finder = FastaiLRFinder()
# To restore the model's and optimizer's states after running the LR Finder
to_save = {"model": model, "optimizer": optimizer}
with lr_finder.attach(trainer, to_save, end_lr=1e-02) as trainer_with_lr_finder:
trainer_with_lr_finder.run(train_loader)
Let’s plot how the learning rate changes within our specified range and print the suggested learning rate.
lr_finder.plot()
print("Suggested LR", lr_finder.lr_suggestion())
Suggested LR 1.0148376909312998e-05
Now we will apply the suggested learning rate to the optimizer, and train the model again with optimal learning rate.
lr_finder.apply_suggested_lr(optimizer)
print(optimizer.param_groups[0]["lr"])
1.0148376909312998e-05
trainer.run(train_loader, max_epochs=3)
State:
iteration: 1407
epoch: 3
epoch_length: 469
max_epochs: 3
output: 0.09644963592290878
batch: <class 'list'>
metrics: <class 'dict'>
dataloader: <class 'torch.utils.data.dataloader.DataLoader'>
seed: <class 'NoneType'>
times: <class 'dict'>
# Calculate the new metrics after using the optimal lr
evaluator.run(test_loader)
print(evaluator.state.metrics)
{'Accuracy': 0.9715, 'Loss': 0.0908882568359375}
As we saw the accuracy increased and loss decreased on the test dataset when we trained our model for the same number of epochs with an optimal learning rate.