Training Models on stages

[bakerhassan]

I would like to train my model in two stages. The first one is to train the model as usual and then at the end of testing, I train a linear model that would be injected between two layers. The second stage is to finetune the model. I found that training a linear model during the testing stage doesn't fit with torch lightning design pattern, so in order to add the necessary information (i.e., the linear model) obtained during testing in the first stage, I created this Callback:

class TestingEndModelCheckpoint(Checkpoint):
    def __init__(self):
        super().__init__()

    def on_test_end(self, trainer, pl_module):
        torchCheckpointIO = TorchCheckpointIO()
        checkpoint_callback: ModelCheckpoint = trainer.checkpoint_callback
        checkpoint = torchCheckpointIO.load_checkpoint(checkpoint_callback.best_model_path)
        if hasattr(pl_module, "best_correction_model"):
            checkpoint['best_correction_model'] = pl_module.best_correction_model.state_dict()
        if hasattr(pl_module, "best_correction_conditional_idx"):
            checkpoint['best_correction_conditional_idx'] = pl_module.best_correction_conditional_idx
        if hasattr(pl_module, "covariate_stage"):
            checkpoint['covariate_stage'] = pl_module.covariate_stage
        torchCheckpointIO.save_checkpoint(checkpoint, checkpoint_callback.best_model_path)

covariate stage is the stage number. Here is how I am planning to run the two stages trainer:

    for stage in range(2):
        modify_callbacks_monitor_parameter(callbacks, stage)

        log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
        trainer: Trainer = hydra.utils.instantiate(cfg.trainer, callbacks=callbacks, logger=logger)

        log.critical(f'Entering stage:{stage}')
        if stage == 1:
            callbacks.append(MNISTFineTuning(model.get_correction_information()[1]))
            ckpt_path = trainer.checkpoint_callback.best_model_path
        if stage == 0:
            ckpt_path = cfg.get("ckpt_path")
        if cfg.get("train"):
            log.info("Starting training!")
            trainer.fit(model=model, datamodule=datamodule, ckpt_path=ckpt_path)

        train_metrics = trainer.callback_metrics

        if cfg.get("test"):
            log.info("Starting testing!")
            ckpt_path = trainer.checkpoint_callback.best_model_path
            if ckpt_path == "":
                log.warning("Best ckpt not found! Using current weights for testing...")
                ckpt_path = None
            trainer.test(model=model, datamodule=datamodule, ckpt_path=ckpt_path)
            log.info(f"Best ckpt path: {ckpt_path}")

My problem now is how to load everything in the second stage to the trainer except itsloops information which controls the training flow. Here is one of my attempts to make the trainer load everything from the first stage and at the same time force it to start training from epoch 0 to max_epoch (i.e., running the usual flow) but it fails as I broke something in the flow:

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        super().on_load_checkpoint(checkpoint)

        # .5 meaning that this is a shift from stage 0 to stage 1 so we neglect the training parameters only during
        # this stage.

        if self.covariate_stage == .5:
            self.best_correction_model.load_state_dict(checkpoint['best_correction_model'])
            self.trainer.fit_loop = _FitLoop(self.trainer, min_epochs=self.trainer.min_epochs,
                                             max_epochs=self.trainer.max_epochs)
            self.trainer.fit_loop.epoch_loop = _TrainingEpochLoop(self.trainer, min_steps=self.trainer.min_steps,
                                                                  max_steps=self.trainer.max_steps)
        self.covariate_stage = 1

Another workaround but I am not sure if it is the right thing is to double the max_epochs per stage:

        if stage > 0:
            cfg.trainer.max_epochs = cfg.trainer.max_epochs * 2

which works fine (meaning the code runs fine).
Any ideas on how to do it?