Video Classification lr_schedulers and custom optimizer

[dudeperf3ct]

What is best way to access length of train_dataloader?

import torch

import flash
from flash.core.data.utils import download_data
from flash.video import VideoClassificationData, VideoClassifier

# 1. Create the DataModule
# Find more datasets at https://pytorchvideo.readthedocs.io/en/latest/data.html
download_data("https://pl-flash-data.s3.amazonaws.com/kinetics.zip", "./data")

datamodule = VideoClassificationData.from_folders(
    train_folder="data/kinetics/train",
    val_folder="data/kinetics/val",
    clip_sampler="uniform",
    clip_duration=1,
    decode_audio=False,
    batch_size=1,
)

# 2. Build the task
model = VideoClassifier(backbone="x3d_xs", 
                        labels=datamodule.labels, 
                        pretrained=True,
                        lr_scheduler=('onecyclelr', {"max_lr" : 1e-3}))

# 3. Create the trainer and finetune the model
trainer = flash.Trainer(
    max_epochs=1, gpus=torch.cuda.device_count(), strategy="ddp" if torch.cuda.device_count() > 1 else None
)
trainer.finetune(model, datamodule=datamodule, strategy="freeze")

Using OneCycleLR requires to specify either total_steps or both epochs and steps_per_epoch. There's a get_num_training_steps() to obtain total steps but how can we use it in above example?

Bug Report:

Traceback (most recent call last):
  File "test_kinetics.py", line 43, in <module>
    trainer.finetune(model, datamodule=datamodule, strategy="freeze")
  File "/opt/conda/lib/python3.8/site-packages/flash/core/trainer.py", line 161, in finetune
    return super().fit(model, train_dataloader, val_dataloaders, datamodule)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 740, in fit
    self._call_and_handle_interrupt(
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1145, in _run
    self.accelerator.setup(self)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu.py", line 46, in setup
    return super().setup(trainer)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 93, in setup
    self.setup_optimizers(trainer)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 354, in setup_optimizers
    optimizers, lr_schedulers, optimizer_frequencies = self.training_type_plugin.init_optimizers(
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 245, in init_optimizers
    return trainer.init_optimizers(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/optimizers.py", line 35, in init_optimizers
    optim_conf = self.call_hook("configure_optimizers", pl_module=pl_module)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1501, in call_hook
    output = model_fx(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/flash/core/model.py", line 500, in configure_optimizers
    return [optimizer], [self._instantiate_lr_scheduler(optimizer)]
  File "/opt/conda/lib/python3.8/site-packages/flash/core/model.py", line 756, in _instantiate_lr_scheduler
    lr_scheduler: Union[_LRScheduler, Dict[str, Any]] = lr_scheduler_fn(optimizer, **lr_scheduler_kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/optim/lr_scheduler.py", line 1174, in __init__
    raise ValueError("You must define either total_steps OR (epochs AND steps_per_epoch)")
ValueError: You must define either total_steps OR (epochs AND steps_per_epoch)

Replacing onecyclelr with hugging face scheduler (installed transformers before running)

lr_scheduler=('cosine_schedule_with_warmup', {"num_warmup_steps" : 0.1})

Bug report:

Traceback (most recent call last):
  File "test_kinetics.py", line 43, in <module>
    trainer.finetune(model, datamodule=datamodule, strategy="freeze")
  File "/opt/conda/lib/python3.8/site-packages/flash/core/trainer.py", line 161, in finetune
    return super().fit(model, train_dataloader, val_dataloaders, datamodule)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 740, in fit
    self._call_and_handle_interrupt(
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1145, in _run
    self.accelerator.setup(self)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/gpu.py", line 46, in setup
    return super().setup(trainer)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 93, in setup
    self.setup_optimizers(trainer)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 354, in setup_optimizers
    optimizers, lr_schedulers, optimizer_frequencies = self.training_type_plugin.init_optimizers(
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 245, in init_optimizers
    return trainer.init_optimizers(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/optimizers.py", line 35, in init_optimizers
    optim_conf = self.call_hook("configure_optimizers", pl_module=pl_module)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1501, in call_hook
    output = model_fx(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/flash/core/model.py", line 500, in configure_optimizers
    return [optimizer], [self._instantiate_lr_scheduler(optimizer)]
  File "/opt/conda/lib/python3.8/site-packages/flash/core/model.py", line 744, in _instantiate_lr_scheduler
    num_training_steps: int = self.get_num_training_steps()
  File "/opt/conda/lib/python3.8/site-packages/flash/core/model.py", line 629, in get_num_training_steps
    dataset_size = len(self.train_dataloader())
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/hooks.py", line 477, in train_dataloader
    raise NotImplementedError("`train_dataloader` must be implemented to be used with the Lightning Trainer")
NotImplementedError: `train_dataloader` must be implemented to be used with the Lightning Trainer

Versions

pl/flash/transformers/torch : 1.5.10 / 0.7.1 / 4.17.0 / 1.9.0

Is it possible to write custom optimizer with different learning rates?

@VideoClassifier.optimizers
def custom_optim():
    return torch.optim.SGD(
        params=[
            {"params": model.head.parameters(), "lr": 4e-3},
            {"params": model.backbone.parameters(), "lr": 4e-4},
        ],
        lr=4e-3,
        momentum=0.9,
        weight_decay=4e-4,
    )
    
# 2. Build the task
model = VideoClassifier(backbone="x3d_xs", 
                        labels=datamodule.labels, 
                        pretrained=False,
                        optimizer='custom_optim')

[karthikrangasai]

Hello @dudeperf3ct ,

The issue NotImplementedError: 'train_dataloader' must be implemented to be used with the Lightning Trainer is because of the breaking change in PL-1.5.x which changes the order of the loading processes and this is being worked up on in the main PL repo.

But until then a fix has to be implemented for it's compatibility with Flash.
CC: @ethanwharris [Please do correct me if I am wrong here]

Regarding the custom optimizers, currently I don't think it is not possible to do something like that.

        optimizers_kwargs: Dict[str, Any] = {}
        optimizer_fn = self._get_optimizer_class_from_registry(self.optimizer.lower())
        
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        optimizer: Optimizer = optimizer_fn(model_parameters, **optimizers_kwargs)

The above lines implement the instantiation of the optimizer inside a flash task. As you can see that the model parameters that are un-frozen will be passed on to the function that is registered. Since a generator of the parameters are directly passed, it is difficult to differentiate between the head or the backbone just using the parameters. So, the answer would be NO. Thanks for pointing this use case out which clearly I had missed out during the implementation. Will fix it soon.

[dudeperf3ct]

@karthikrangasai Any solution to one cycle lr about getting total_steps or steps_per_epoch?

[karthikrangasai]

You can take help from this PR #11599. It has the correct implementation that gets the estimated number of training steps. This could of some help in the case you would want to subclass Trainer class and implement one.