Resuming should allow to differentiate what to resume (steps/opti/weights)

[thoglu]

Currently it is possible to either resume only the full training state (epoch/global steps / optimizer / scheduler options / and weights), or only the weights.

I would like to be able to switch the optimizer at some point, i.e. skip restoring optimizer/scheduler, but still load the epoch/global steps. I only see a way to do this with hacks at the moment. Any other way? Could this be a feature, to specify in the trainer.init function specifically what to restore/what not to restore?

[rohitgr7]

one way this can be done is by updating the error to just a warning and make relevant changes below:
https://github.com/PyTorchLightning/pytorch-lightning/blob/d20fd8e5ab1a52747fee2cd53290a679d8b726d0/pytorch_lightning/trainer/connectors/checkpoint_connector.py#L120-L124

cc: @PyTorchLightning/core-contributors thoughts?

[tchaton]

Hey @rohitgr7,

I am not sure resuming and changing optimizer / scheduler is the best option.

Possibly, we could extend LightningModule configure_optimizers to take current_epoch.

def configure_optimizers(self, current_epoch: int = None):


     return [optimizers], [lr_schedulers], should_update: bool

Internally, we inspect the model function. If the user requested current_epoch, we call this function on every_epoch and change its optimizers, lr_schedulers only if should_update = True

Example:

def configure_optimizers(self, current_epoch: int = None):
    milestones = [0, 100, ...]
    optimizers = None
    lr_schedulers = None
    if current_epoch == milestones[0]:
        # init optimizers / schedulers 

    elif current_epoch == milestones[-1]:
        # init new optimizers / schedulers 

    ...

    return [optimizers], [lr_schedulers], current_epoch  in current_epoch milestones

What are your thoughts ?

[rohitgr7]

@tchaton yeah this is a good suggestion, but what if someone sets save_weights_only in ModelCheckpoint and uses resume_from_checkpoint using that checkpoint? In such a case this is still a problem.

[thoglu]

I brought this up because the current options (save_weights_only or not) are a little too restrictive. I have found a hack for now for me but there is I think a more widespread use-case to be able to change the optimizers, as SGD for example has certain properties that Adam has not etc., and one might want to exploit those at different stages of training.

[carmocca]

What are your thoughts ?

I think that's too complicated and its complexity will grow too large if we wanted to do it in different occasions other than epoch

In general, I think we should let users save any combination of {training, optimizer, model} state and resume training with whichever combination is provided.

[thoglu]

Is there any progress how/when such a combinatorial choice for the users will be implemented?

[carmocca]

No current progress/plans that I know of. Implementing this will require attention