How to use LightningCLI to resume training from checkpoint

[KirillShmilovich]

I'm trying to use the LightningCLI to resume training from checkpoint, but don't understand how to pass keyword arguments of the checkpoint path to enable restoring the model and trainer states from the checkpoint, my current workaround is the following:

def cli_main():
    class MyCLI(LightningCLI):
        def add_arguments_to_parser(self, parser):
            parser.link_arguments("data.unmask_index", "model.unmask_index", apply_on="instantiate")

    ckpt_path = 'my/checkpoint/path/ckpt.ckpt'
    cli = MyCLI(
        DDPM,
        InMemoryDataModule,
        save_config_overwrite=True,
        seed_everything_default=4321,
        run=False,  # used to de-activate automatic fitting.
        trainer_defaults={
            "accelerator": "gpu",
            "devices": 1,
            "max_epochs": 2500000,
            "enable_progress_bar": False,
            #"limit_val_batches": 0.25,
            "check_val_every_n_epoch": 30,
            # "num_sanity_val_steps": 0,
            # "logger": TensorBoardLogger(save_dir="./",name="DDPM_timesteps"),
        },
    )
    cli.trainer.fit(cli.model.load_from_checkpoint(ckpt_path), datamodule=cli.datamodule, ckpt_path=ckpt_path)

note that I call cli.trainer.fit(cli.model.load_from_checkpoint(ckpt_path), datamodule=cli.datamodule, ckpt_path=ckpt_path) specifying the checkpoint path.

Is there a better/recommended method for dealing with restarting from checkpoints with the LightningCLI?

[tshu-w]

Just python cli.py fit --ckpt_path PATH, you can python cli.py fit --help to check all available options.

[adamjstewart]

Is there a way to automatically select the most recent checkpoint or best checkpoint? I would like to do this but don't always know the filepath to my checkpoint file.

[adamjstewart]

--ckpt_path last doesn't seem to work for me, although it's in the help message.

[KirillShmilovich]

Then only option involving a checkpoint appears to be in the trainer.resume_from_checkpoint option that appears to now be deprecated -- this also appears to not address the issue of reinstating the model state & weights from the checkpoint, because this option presumably only impacts the trainer state?

[mauvilsa]

I would say yes, considering the description of ckpt_path: "Path/URL of the checkpoint from which training is resumed".

If you run LightningCLI with run=True and fit --ckpt_path=path internally what is run is cli.trainer.fit(model=cli.model, datamodule=cli.datamodule, ckpt_path=path). I would guess that the .load_from_checkpoint(ckpt_path) in the code of the question is unnecessary.

[KirillShmilovich]

This worked! Thanks!

[adosar]

@mauvilsa If we call fit --config=config.yaml --ckpt_path=path and config.yaml configures optimizer, does this interferes in any way with the state of the optimizer saved in path (due to checkpointing)? Or path just load the LightningModule’s state_dict? This might be the case if one wants to resume training but with a different optimizer.

[mauvilsa]

@adosar I am only aware about how the CLI works, and it just provides parameters and instantiates classes. What you want to know depends on what fit does internally when it receives a ckpt_path argument. I am not sure if this is clearly explained in the docs or you will have to look at the code or experiment to understand what exactly happens.

[adosar]

@mauvilsa Yeap, I would need to look how the trainer works under the hood.

[KasuganoLove]

Hello! There is a similar question. I have already trained an encoder saved in encoder_ckpt.ckpt. I want to train another autoencoder using encoder_ckpt.ckpt as the frozen encoder. How should I configure the LightningCLI?