How to redefine optimizer for pl.LightningModule?

[Godofnothing]

The great advantage and convenient feature of PyTorch-Lightning over vanilla Pytorch is the more convenient interface for optimization loops. Instead of reinventing the wheel every time with manual train and test function, one can just define
training_step and validation_step, e.t.c.

The setup of the optimizer is done in the configure_optimizers method https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers, where one can create optimizers and schedulers.

However, it seems, that this choice fixes the optimizer properties in the definition of the class, and often you would like to try several optimizers and change their params. One way to do it - is to pass additional arguments to the constructor of
pl.LightningModule to be processed in configure_optimizers, where an appropriate optimizer will be chosen.

Also one may try to access directly the optimizers outside the class, but probably it is an unsafe approach.

Is there some alternative way to reconfigure optimizer, like in Tensorflow, where models have complie method https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile?

[justusschock]

Hi @Godofnothing ,

I am afraid there is no clean way to do so. What I personally do to hack around this (if I need it) is that I have arguments to the model that define what is to be returned by configure_optimizers and then call self.trainer.accelerator.setup_optimizers(self.trainer) in any of the other hooks. That makes sure configure_optimizers is called again and converts everything to the correct types. Inside your configure_optimizers function you would have to handle the creation of the optimizers for the different stages in training then.

[tchaton]

Hey @Godofnothing,

If you are using the LightningCLI, this would be simplified slightly.

See this PR: #8815

Best,
T.C

[alvaro-budria]

Hi @justusschock, I am calling self.trainer.accelerator.setup_optimizers(self.trainer), which raises this error:
AttributeError: 'CUDAAccelerator' object has no attribute 'setup_optimizers'.

I am using version 1.7.7 of pytorch-lightning.

Did the interface change since your original post?

[alvaro-budria]

Turns out that by calling self.trainer.strategy.setup_optimizers(self.trainer) it works.

[adosar]

Not the cleanest solution, but you can also overwrite the configure_optimizers method.

For example, assuming you have a custom LightningModule named Foo:

# Define a custom function with the same signature as configure_optimizers.
def custom_config_optimizers(self):
    ...
    return optimizer

Then assign it to your class:

Foo.configure_optimizers = custom_config_optimizers

[Godofnothing]

I've come up with this kind of solution (dangerous and breaking the law, but superficially woking one). Add this two methods to the pl.LightningModule:

@staticmethod
def _validate_compile_args(args):

    def _check_single(args):
        if isinstance(args, torch.optim.Optimizer):
            return True
        return False

    def _check_optimizer_list(args):
        if isinstance(args, list):
            for arg in args:
                if not _check_single(arg):
                    return False
        return True

    def _check_dict(args):
        if isinstance(args, dict):
            assert args.get("optimizer"), "optimizer has to be specified"
            assert isinstance(args["optimizer"], torch.optim.optimizer.Optimizer)
        return True

    def _check_tuple_of_dicts(args):
        if isinstance(args, tuple):
            for arg in args:
                if not _check_dict(arg):
                    return False
        return True

    def _check_none(args):
        if not args:
            return True

    if _check_single(args):
        return
    if _check_optimizer_list(args):
        return
    if _check_dict(args):
        return
    if _check_tuple_of_dicts(args):
        return
    if _check_none(args):
        return

    raise ValueError("Unsupported configuration of optimizers")

def compile(self, *args):
    self._validate_compile_args(args)

    def _configure_optimizers(self):
        return args

    self.configure_optimizers = MethodType(_configure_optimizers, self)
    self.configure_optimizers()

And the example of the use case:

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max')

model.compile({
    "optimizer" : optimizer,
    "lr_scheduler" : {
        "scheduler" : scheduler,
        "monitor" : "val/loss"
    }
})