Fix lr_finder gradient calculation for mode="exponential"

[SkafteNicki]

What does this PR do?

Fixes #14248

Fix the gradient calculation for lr_finder when suggestion a learning rate to start from. Currently the gradient when mode=exponential does not account for the non-homogenous stepping size which leads to wrong approximate of learning rate.

An example, on master this is the current suggested lr

with the fix

You can see that in the second plot, the red point is located earlier on the downward slope — closer to the point where the curve transitions from flat to rapidly decreasing. In contrast, the first plot’s red point is further down the slope, where the loss is still decreasing but less sharply.

Expand for script

#!/usr/bin/env python3
import lightning as L
import torch
import torchvision
import os


# Set seed for reproducibility
L.seed_everything(42)


class SimpleModel(L.LightningModule):
    def __init__(self, lr=1e-4):
        super().__init__()
        self.lr = lr
        self.model = torch.nn.Sequential(
            torch.nn.Flatten(),
            torch.nn.Linear(28 * 28, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, 10)
        )

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

    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 = "MNIST_test/", batch_size: int = 64):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size

    def setup(self, stage: str):
        os.makedirs(self.data_dir, exist_ok=True)
        
        mnist_full = torchvision.datasets.MNIST(
            self.data_dir, train=True, download=True, 
            transform=torchvision.transforms.ToTensor()
        )
        self.mnist_train, _ = torch.utils.data.random_split(mnist_full, [5000, 55000])

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


def main():
    L.seed_everything(42)

    # Create model and data
    model = SimpleModel()
    dm = MNISTDataModule()
    
    
    from lightning.pytorch.callbacks import LearningRateFinder
    lr_finder = LearningRateFinder(
        min_lr=1e-6,
        max_lr=1,
        num_training_steps=100,
        mode='exponential',
        early_stop_threshold=4.0,
        update_attr=True,
    )

    trainer = L.Trainer(
        max_epochs=1,
        callbacks=[lr_finder],
        enable_progress_bar=False,
        deterministic=True
    )
    trainer.fit(model, dm)

    lr_finder.optimal_lr.plot(suggest=True, show=True)
    print(f"Suggested learning rate: {lr_finder.optimal_lr.suggestion}")


if __name__ == "__main__":
    main()

``

</details>



<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [ ] Did you read the [contributor guideline](https://github.com/Lightning-AI/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [ ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


<!-- readthedocs-preview pytorch-lightning start -->
----
📚 Documentation preview 📚: https://pytorch-lightning--21055.org.readthedocs.build/en/21055/

<!-- readthedocs-preview pytorch-lightning end -->