How to utilize timm's scheduler?

[sooperset]

🚀 Feature

Hi, I want to reproduce a result of image classification network by using timm library.
But I couldn't use timm.scheduler.create_scheduler because pytorch_lightning doesn't accept custom class for a scheduler.
(timm.scheduler is not the torch.optim.lr_scheduler class)

from timm.scheduler import create_scheduler
from timm.optim import create_optimizer

def configure_optimizers(self):
    optimizer = create_optimizer(self.args, self)
    scheduler, _ = create_scheduler(self.args, optimizer)

    return {
        'optimizer': optimizer,
        'lr_scheduler': scheduler,
    }

Then results this error

File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/optimizers.py", line 141, in configure_schedulers
raise ValueError(f'The provided lr scheduler "{scheduler}" is invalid')

Is there a plan for utilizing timm's scheduler?

Motivation

Pitch

Alternatives

Additional context

[DuinoDu]

One solution is to wrap it into a callback. I met similiar problem, and I created a MMCVLrCallback. Hope help.

[tchaton]

Dear @soomiles,

It shouldn't be hard to extend the conditions to support it.
Feel free to make a PR.

Best,
T.C

[potipot]

Hello, has there been any update on this? @soomiles, have you perhaps worked it out?

[DuinoDu]

Dear @soomiles,

It shouldn't be hard to extend the conditions to support it.
Feel free to make a PR.

Best,
T.C

I plan to make a PR in these days.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[jinwon-samsung]

Is the timm scheduler still not supported?

[dgmp88]

This wont work with multiple parameter groups, but a quick fix if anyone else want to use it:

from timm.scheduler import StepLRScheduler
class TimmStepLRScheduler(torch.optim.lr_scheduler.LambdaLR):
    def __init__(self, optim, **kwargs):
        self.init_lr = optim.param_groups[0]["lr"]
        self.timmsteplr = StepLRScheduler(optim, **kwargs)
        super().__init__(optim, self)

    def __call__(self, epoch):
        desired_lr = self.timmsteplr.get_epoch_values(epoch)[0]
        mult = desired_lr / self.init_lr
        return mult

[timothylimyl]

Issue still exists for me (updated to latest lightning version) upon running trainer.fit(model, data):

line 505, in _update_learning_rates
    lr_scheduler["scheduler"].step()
TypeError: step() missing 1 required positional argument: 'epoch'

The fix according to the PR that I followed:

    def configure_optimizers(self):

        return [self.optimizer], [{"scheduler": self.scheduler, "interval": "epoch"}]

    def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
        print(self.current_epoch)
        scheduler.step(
            epoch=self.current_epoch
        )  # timm's scheduler need the epoch value

The scheduler object was: <timm.scheduler.cosine_lr.CosineLRScheduler object at 0x7f6d1dc29be0>. When debugging, it does not seem to go to the function lr_scheduler_step. It will run successfully for a single training epoch and then the error above was raised (which makes sense since LR is scheduled every epoch), this prevents going to the validation set.

[timothylimyl]

Upon further debugging, I realised that the latest lightning (1.5.10) is not updated yet with this merge. Silly me, I was assuming that the latest tag will follow the master branch as the practice is to only merge when stable and necessary. I saw that it was merged to master and immediately assume that it will be released upon the new updates, but further reading in details (looking at dates after this merge), I see that it was not brought up in any updates. Later on, I found out that it is actually going to only be tagged in the upcoming 1.6 version.

Sharing this for learning purposes in case someone was inexperienced as I am. Remember to check the PR milestone!