Learning rate finder using Lightning CLI

[theja-vanka]

How to trigger lr_finder using Lightning CLI ?

[marianne-omnius]

I implemented it the following way:

my_lightning_module.py

class MyLightningModule(LightningModule):
    def __init__(self):
        super().__init__()
        self.lr = None  # Gets set by LearningRateFinder if configured

    def on_fit_start(self):
        if self.lr:
            # Overwrite learning rate after running LearningRateFinder
            for optimizer in self.trainer.optimizers:
                for param_group in optimizer.param_groups:
                    param_group["lr"] = self.lr

Then I add the LearningRateFinder callback to my training config as follows:

config.yaml

trainer:
  callbacks:
    - class_path: pytorch_lightning.callbacks.LearningRateFinder

Then I run the training using LightningCLI and passing the above config:

python lightning_cli.py fit --config config.yaml

Hope that helps! 🙂

[theja-vanka]

Thanks you for your help. It works like a charm.

[theja-vanka]

Hey,
I tried to play around with my module and was unable to use it. I have enclosed my snippet. Kindly guide where I went wrong.

import os
import torch
import mlflow
from torch import nn
import lightning.pytorch as L
from watermark import watermark
from lightning.pytorch.cli import LightningCLI
from torchmetrics.classification import BinaryAccuracy
from lightning.pytorch.callbacks import RichProgressBar
from lightning.pytorch.callbacks import ModelCheckpoint
from torch.optim.lr_scheduler import CosineAnnealingLR

from src.models import TemporalModel
from src.data import TemporalDataModule

os.environ["LOGNAME"] = "krishnatheja.vanka"


class Config:
    CRITERION = (
        nn.BCEWithLogitsLoss()
    )  # nn.BCEWithLogitsLoss(pos_weight=torch.tensor([8]))
    OPTIMIZER = torch.optim.AdamW


class LightningModel(L.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.save_hyperparameters()

        self.lr = None
        self.config = config
        self.criterion = Config.CRITERION
        self.optimizer = Config.OPTIMIZER

        self.model = TemporalModel(config)
        self.model = torch.compile(self.model)

        self.train_acc = BinaryAccuracy()
        self.val_acc = BinaryAccuracy()
        self.test_acc = BinaryAccuracy()

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

    def on_fit_start(self):
        if self.lr:
            # Overwrite learning rate after running LearningRateFinder
            for optimizer in self.trainer.optimizers:
                for param_group in optimizer.param_groups:
                    param_group["lr"] = self.lr

    def _shared_step(self, batch):
        sample, target, hidden = batch
        logits, hidden = self(sample, hidden)
        predicted = torch.squeeze(logits[-1], dim=0)
        loss = self.criterion(predicted, target.float())
        return loss, target, predicted

    def training_step(self, batch, batch_idx):
        loss, true_labels, predicted_labels = self._shared_step(batch)

        predicted_labels = torch.sigmoid(predicted_labels)

        self.log("train_loss", loss)
        self.train_acc(predicted_labels, true_labels)
        self.log(
            "train_acc", self.train_acc, prog_bar=True, on_epoch=True, on_step=False
        )
        return loss

    def validation_step(self, batch, batch_idx):
        loss, true_labels, predicted_labels = self._shared_step(batch)
        predicted_labels = torch.sigmoid(predicted_labels)
        self.log("val_loss", loss, prog_bar=True)
        self.val_acc(predicted_labels, true_labels)
        self.log("val_acc", self.val_acc, prog_bar=True)

    def test_step(self, batch, batch_idx):
        _, true_labels, predicted_labels = self._shared_step(batch)
        predicted_labels = torch.sigmoid(predicted_labels)
        self.test_acc(predicted_labels, true_labels)
        self.log("test_acc", self.test_acc)

    def configure_optimizers(self):
        optimizer = self.optimizer(self.parameters(), lr=self.lr)

        return optimizer


def cli_main():
    checkpoint_callback = ModelCheckpoint(
        monitor="val_loss",
        dirpath="artifacts/fps-5_iter_6/",
        filename="point-{epoch:02d}-{val_loss:.2f}",
        mode="min",
        save_top_k=5,
    )
    mlflow.set_experiment(experiment_name="yawn-temporal")
    mlflow.pytorch.autolog()
    cli = LightningCLI(
        LightningModel,
        TemporalDataModule,
        seed_everything_default=42,
        run=False,
        trainer_defaults={
            "callbacks": [checkpoint_callback, RichProgressBar()],
        },
    )

    cli.trainer.fit(cli.model, datamodule=cli.datamodule)
    cli.trainer.test(ckpt_path="best", datamodule=cli.datamodule)


if __name__ == "__main__":
    # python main.py --trainer conf/trainer.yaml --model conf/model.yaml --data conf/data.yaml
    print(watermark(packages="torch,lightning", python=True))
    print("Torch CUDA available?", torch.cuda.is_available())
    cli_main()

My config.yaml file is

accelerator: auto
strategy: auto
devices: auto
num_nodes: 1
enable_model_summary: true
max_epochs: 300
deterministic: true
callbacks:
  class_path: lightning.pytorch.callbacks.LearningRateFinder
logger:
  class_path: lightning.pytorch.loggers.TensorBoardLogger
  init_args:
    save_dir: logs/
    name: iter_6_fps5

[mauvilsa]

You give some code and yaml, but didn't really explain what went wrong. A bit difficult to help without this.

A couple details based on what I can see in the code. Though I have no idea if this has any relation to your problem.

• It is not recommended to use class instances as defaults like done here in trainer_defaults. Either give a dict with class_path and init_args or use lazy_instance. Have a read at class-type-defaults.
• The idea of using LightningCLI is that modules define in their __init__ parameters with good names and with type hints. If done as expected, the parameters become configurable from command line and yaml and there are nice features like --print_config. A single config parameter goes against this. For example have a look at how to make the optimizer configurable in multiple-optimizers-and-schedulers.

And one more tip. When there are problems related to the CLI, see how-do-i-troubleshoot-a-cli.

[theja-vanka]

Hi @mauvilsa ,

Thank you for you feedback. I hid the model class instance as this is related to work. But I will definitely follow your advice moving forward.

The issue that I am facing is that, I am trying to use learning rate finder in order to minimize guess work and leverage its capabilities. However, the documentation is a bit dry and does not provide a straight forward way of integrating it (unlike the other features).

I will try to experiment with how-do-i-troubleshoot-a-cli and see if I can resolve the query.

Thanks again for your advice and any help here would be great.

[marianne-omnius]

Two more comments from looking at your code / config:

1. I think configure_optimizers gets called before on_fit_start. So you need an initial learning rate (e.g. 0.001) when you first create the optimizer in configure_optimizers. Learning rate finder will then overwrite it later. In case learning rate finder is not configured or learning rate finder cannot determine a good learning rate, then your initial default learning rate will be used for training.
2. In your config, callbacks should be a list. So you need a - before class_path to make it a list in yaml (see my example config).

In general I agree that the usage is not really clear from the documentation. So all my comments here are my best guess from trying to make it work for myself, not sure if it is the "right" way of using it.