How do I use the LR Scheduler on a general task training?

[jin0g]

Hi,
This project is sophisticated, exciting for me.

How do I use the LR Scheduler on a general task training?
As far as I saw, flash.Task is not designed to allow me to configure the LR scheduler.
Also, there is no way to set any parameters other than the LR on the optimizer (like momentum, weight_decay, etc.)

To get around this, I created a class inherited from flash.Task as follows:

class MyTask(flash.Task):
    def __init__(self, model, optimizer: torch.optim.Optimizer, scheduler: dict, *args, **kwargs):
        super().__init__(model, *args, **kwargs)
        self.optimizer = optimizer
        self.scheduler = scheduler
        
    def configure_optimizers(self):
        return [self.optimizer], [self.scheduler]

and use like:

model = ...
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4)
scheduler = {'scheduler': CosineAnnealingLR(optimizer, T_max=200), 'interval': 'epoch', 'frequency': 1}

task = MyTask(model, loss_fn=torch.nn.functional.cross_entropy, optimizer=optimizer, scheduler=scheduler)

flash.Trainer(...).fit(task, train_data_loader, val_data_loader)

(Reason why I wanted to adopt this was to get reproduction: https://github.com/kuangliu/pytorch-cifar)

Actually, I don't like this extension.
This loses the simplicity of the flash.

Is there any good ideas? Thank you.

[aribornstein]

I believe, a scheduler should be exposed to the tasks as an optional argument in https://github.com/PyTorchLightning/lightning-flash/blob/master/flash/core/model.py. That would enable this to be much cleaner.