Scripting a model defined in a LightningCLI yaml

[nicolas-busca]

Hi Lightning!

I'm looking for a way to script a model defined via a LightningCLI yaml.
I could parse the yaml and instantiate all the classes, etc, but
since lightning is able to load a model from a checkpoint via --ckpt_path,
all the infrastructure must be there already.

However, --ckpt_path isn't available when run=False in LightningCLI, so
I am unable to load the model's weights.

What is the recommended path forward?

Thanks!

[nicolas-busca]

For what is worth, I'm doing:

cli = LigthnignCLI(run=False)

ckpt = torch.load(ckpt_path)
cli.model.load_state_dict(ckpt['state_dict'])

but it feels a bit contrived...

[mauvilsa]

If you train a model that does self.save_hyperparameters(), then it should be possible to use load_from_checkpoint as explained in lightning_cli_advanced_3.html#models-with-multiple-submodules. The only caveat is this bug #20311, for which there is already a pull request with a fix, but not yet merged.

[nicolas-busca]

Thanks Maricio for the quick reply.

I don't have self.save_hyperparameters() in my class, how is LightningCLI able to instantiate the model when I continue training from a previous checkpoint via --ckpt_path ... ? Why is self.save_hyperparameters() needed with run=False?

Are there any caveats with the workaround I shared above that I might not be aware of?

[mauvilsa]

With self.save_hyperparameters() the parameters used to instantiate the model during training are included in the checkpoint. If the parameters are not in the checkpoint, then it is impossible to instantiate with the parameters used before. It is up to you to provide the needed parameters again.

Take for example a model that receives number_of_classes, which would be given a particular value to fit. If you manually load_state_dict as above, then you need to remember to provide the correct number_of_classes when instantiating LigthnignCLI. If you don't provide number_of_classes then model instantiation would fail either because number_of_classes not given, or load_state_dict fails because tensor shapes don't match due to a wrong number_of_classes. This applies to all parameters that the model receives.

Could be that in your particular case the model doesn't have parameters, or you actually want to load_state_dict using the default parameters of the model instead of what was used for fit.

[nicolas-busca]

I understand, that would be the case without the yaml config. But typically I have a yaml that specifies all the model parameters (e.g. number of classes) that I use to train a model. That gives me a checkpoint that I would like to use together with the yaml to instantiate the model (effectively like having the --ckpt_path option available when run=False) so I don't need to reimplement parsing the yaml and going down the injection tree...