Predicting from checkpoint with CLI

[Northo]

I want to predict with a model for which I have a checkpoint. The arguments used to train the model (producing the checkpoint), causes the model to have different layer sizes than the default values. When using ckpt_path, this then causes the loading to fail, due to a size mismatch.

# cli.py
from lightning.pytorch.cli import LightningCLI
from lightning.pytorch.demos.boring_classes import DemoModel, BoringDataModule

def cli_main():
    cli = LightningCLI(DemoModel, BoringDataModule)


if __name__ == "__main__":
    cli_main()

python cli.py fit --trainer.max_epochs 1 --model.out_dim 5
python cli.py predict --ckpt_path <path_from_previous_step>

RuntimeError: Error(s) in loading state_dict for DemoModel:
        size mismatch for l1.weight: copying a param with shape torch.Size([5, 32]) from checkpoint, the shape in current model is torch.Size([10, 32]).
        size mismatch for l1.bias: copying a param with shape torch.Size([5]) from checkpoint, the shape in current model is torch.Size([10]).                                                                   ─╯

Two obvious workarounds are

1. explicitly giving the appropriate flags in the predict call (--model.some_layer_size 4 above), or
2. rewriting the CLI to use the LightningModule.load_from_checkpoint.

Explicitly giving all the flags becomes very cumbersome when there are many. The second approach might be a good one, though I haven't yet tried implementing it.

Is there something I am overlooking here? What is the recommended way to predict from a checkpoint like this?

Any input greatly appreciated:)

[Northo]

After giving this some more thought, I suppose the generated cli_config.yaml file from fit could be passed to predict. However, that only applies in this simplified setup. For example, the situation that caused me to raise this question, is one where we have performed a hyperparameter sweep (without the Lightning CLI), and want to use the checkpoints to predict.

[Northo]

My current workaround is

class MyLightningCLI(LightningCLI):
    def after_instantiate_classes(self) -> None:
        if self.subcommand == "predict" and self.config["predict"]["ckpt_path"]:
            logger.warning("Hijacking model after instantiate.")
            config_from_cli = self.parser.parse_args(self.parser.args, defaults=False)
            model_args = config_from_cli["predict"].get("model", {})

            logger.debug(f"{model_args=}")
            self.model = MyModel.load_from_checkpoint(
                self.config["predict"]["ckpt_path"], **model_args
            )
            logger.debug(f"Model loaded from checkpoint: {self.model}")

This is not great, but works for my current use-case. I suppose the self.subcommand == "predict" check could be removed to work with other commands as well, but that is currently not something I need.