lr_find suggestion ignores lr steps

[kotchin]

🐛 Bug

lr_find computes suggestions using np.gradient() but only provides the array of losses, ignoring the lr at which each loss was computed. As such, np.gradient() assumes homogenous stepping size, which is not true when using exponential mode, which is the default mode.
This is visible here: https://github.com/Lightning-AI/lightning/blob/acd4805f1a284e513272d150de6f98f27a0489b3/src/pytorch_lightning/tuner/lr_finder.py#L201
Which is executed irrespective of the mode.

Possible fix of suggestion() in lr_finder.py (I'm not addressing the already present TODO):

losses = np.array(self.results["loss"][skip_begin:-skip_end])
losses = losses[np.isfinite(losses)]
lrs = np.array(self.results["lr"][skip_begin:-skip_end])
lrs =  losses[np.isfinite(losses)]

if len(losses) < 2:
    # computing np.gradient requires at least 2 points
    log.error(
        "Failed to compute suggestion for learning rate because there are not enough points. Increase the loop"
        " iteration limits or the size of your dataset/dataloader."
    )
    self._optimal_idx = None
    return None

# TODO: When computing the argmin here, and some losses are non-finite, the expected indices could be
#   incorrectly shifted by an offset
# Kotchin's note: simulatneous attempt at addressing the TODO
min_grad = np.gradient(losses, lrs).argmin()
self._optimal_idx = np.arange(len(self.results["loss"]))[skip_begin:-skip_end][np.isfinite(losses)][min_grad]
return self.results["lr"][self._optimal_idx]

To Reproduce

It's trivial, but here's a minimal code to highlight how this matters:

import matplotlib.pyplot as plt
import numpy as np

losses = np.sin(np.arange(0, 6, 0.1))
lrs = np.exp(np.arange(-9, -1, (-1 - -9)/len(losses)))

plt.plot(lrs, losses)
min_grad_pl = np.gradient(losses).argmin() # Current PL implementation
plt.scatter(lrs[min_grad_pl], losses[min_grad_pl], c='red', label='PL suggestion')
min_grad = np.gradient(losses, lrs).argmin() # Expected PL implementation
plt.scatter(lrs[min_grad], losses[min_grad], c='green', label='expected suggestion')
plt.xscale('log')
plt.legend()
plt.xlabel('lr')
plt.ylabel('loss')
plt.show()

Executing the code below will yield the follow illustration which should highlight the issue:

Expected behavior

Yield the correct lr suggestion, accounting for the non-homogeneous lr steps.

Environment

pytorch-lightning 1.7.1
python 3.10.4

Additional context

EDIT: added code suggestion placement detail.
EDIT2: Attempt at addressing the TODO in the suggestion() function.

cc @akihironitta @Borda @rohitgr7

[kotchin]

I'm facing some github issues so I can't actually make a PR myself, but I wanted to share the below code which I'd propose to make it work, addressing the commented TODO in the passing, hopefully in a satisfactory way. This is intended as a drop-in replacement for the suggestion() function:

    def suggestion(self, skip_begin: int = 10, skip_end: int = 1) -> Optional[float]:
        """This will propose a suggestion for an initial learning rate based on the point with the steepest
        negative gradient.

        Args:
            skip_begin: how many samples to skip in the beginning; helps to avoid too naive estimates
            skip_end: how many samples to skip in the end; helps to avoid too optimistic estimates

        Returns:
            The suggested initial learning rate to use, or `None` if a suggestion is not possible due to too few
            loss samples.
        """
        losses = np.array(self.results["loss"][skip_begin:-skip_end])
        lrs = np.array(self.results["lr"][skip_begin:-skip_end])
        is_finite = np.isfinite(losses)
        losses = losses[is_finite]
        lrs =  lrs[is_finite]

        if len(losses) < 2:
            # computing np.gradient requires at least 2 points
            log.error(
                "Failed to compute suggestion for learning rate because there are not enough points. Increase the loop"
                " iteration limits or the size of your dataset/dataloader."
            )
            self._optimal_idx = None
            return None

        min_grad = np.gradient(losses, lrs).argmin()
        all_losses_idx = np.arange(len(self.results["loss"]))
        idx_non_skipped = all_losses_idx[skip_begin:-skip_end]
        idx_finite = idx_non_skipped[is_finite]
        self._optimal_idx = idx_finite[min_grad]
        return self.results["lr"][self._optimal_idx]

This has been validated on the following snippet of code:

import lightning as L
import torch, torchvision
import matplotlib.pyplot as plt

#https://pytorch-lightning.readthedocs.io/en/stable/advanced/training_tricks.html#learning-rate-finder

L.seed_everything(42)

class LitModel(L.LightningModule):
    def __init__(self, lr=1e-4):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.lr = lr

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = torch.nn.functional.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr)

class MNISTDataModule(L.LightningDataModule):
    def __init__(self, data_dir: str = "path/to/dir", batch_size: int = 32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: str):
        self.mnist_test = torchvision.datasets.MNIST(self.data_dir, train=False, download=True, transform=torchvision.transforms.ToTensor())
        self.mnist_predict = torchvision.datasets.MNIST(self.data_dir, train=False, download=True, transform=torchvision.transforms.ToTensor())
        mnist_full = torchvision.datasets.MNIST(self.data_dir, train=True, download=True, transform=torchvision.transforms.ToTensor())
        self.mnist_train, self.mnist_val = torch.utils.data.random_split(mnist_full, [55000, 5000])

    def train_dataloader(self):
        return torch.utils.data.DataLoader(self.mnist_train, batch_size=self.batch_size)

boring_model = LitModel()
boring_dm = MNISTDataModule(data_dir='MNIST_tmp/')

trainer = L.Trainer(   logger=False, 
                        enable_checkpointing=False,
                        auto_lr_find=True)
lr_finder = trainer.tuner.lr_find(boring_model, boring_dm)

# Results can be found in
print("LR finder results:", lr_finder.results)
# Pick point based on plot, or get suggestion
new_lr = lr_finder.suggestion()
print("Found LR:", new_lr)

plt.plot(lr_finder.results['lr'], lr_finder.results['loss'], color='blue', label='lr finder loss')
#plt.xscale('log') # Can comment out leave in to compare the results as it can be difficult to see optimal lr in log scale
ymin, ymax = plt.ylim()
plt.vlines(new_lr, ymin, ymax, color='red', label='suggestion')
plt.legend()
plt.show()

Previously, optimal lr was 0.09120108393559097 and now optimal lr is 0.0007585775750291836. 2 orders of magnitude of difference.

This was tested on:
python=3.10
torch=1.12
torchvision=0.13
lightning=1.9.0.dev0

Edit: removed an unnecessary import from the test code.
Edit2: also, in the following link it is suggested to plot the lr_finder results using fig = lr_finder.plot(suggest=True) then fig.show(), but this only results in a brief figure flashing, then instantly disappearing, preventing from viewing. This can be fixed using fig = lr_finder.plot(suggest=True, show=True) but the code or the documentation should either/both be fixed.

[mrembalski]

@Borda this seems like a drop-in replacement. I can open a PR with this code if @kotchin has issues/does not want to do that anymore. What do you think?

[kotchin]

@mrembalski indeed, I'm still facing the same issue on my account and have had to halt all MR... If anyone else is willing to merge, it would be helpful.