Evaluate or change type from argument in configuration yaml file. Fix Jsonargparse error "'cannot represent an object'"

[LysSanzMoreta]

Hi!

I am new to the Pytorch Lighning CLI and I am struggling modifying an existing pipeline to be able to load checkpoints into the model.

At the moment, in the set up we are using pl.LightningModule, pl.LightningDataModule and a config.yaml file.

So far, I have been able to correctly initialize/instantiate the LighningModule by adding a CheckPointLoader constructor to the Yaml SafeLoader class.

.....
model:
  model:
    class_path: Model
    init_args:
      precomputed_checkpoint:
         - !CheckpointLoader 
           file_path: ../notebooks/lightning_logs/version_7/checkpoints/epoch=49-step=3150.ckpt
           attr: ...layer1.bias # I extract the torch.Tensor "bias" parameter from the layer1 from the checkpointed model
           convert_fn: ...transforms.load.process_checkpoint #returns a torch.Tensor
.....

The CheckpointLoader uses the load_from_checkpoint function to retrieve the parameter bias from the checkpointed model (or whichever parameter of choice). This means that it returns a list with the torch.Tensor (Side note:I am not sure yet at which point the torch.Tensor returned by the convert_fn is wrapped into a list and I would also like to avoid that).

After the instantiation of pytorch LighningModule class, then the cli proceeds to parse the command line arguments cli(args=args). Then it struggles because the CheckPointLoader returned a list(torch.Tensor) which cannot be interpreted or evaluated by the parser.

  File "/opt/conda/lib/python3.10/site-packages/jsonargparse/_typehints.py", line 610, in raise_union_unexpected_value
    raise ValueError(
ValueError: Does not validate against any of the Union subtypes
Subtypes: (<class 'Model'>, <class 'NoneType'>)
Errors:
  - ('cannot represent an object', tensor(..., device='meta', size=(16,)))
  - Expected a <class 'NoneType'>
Given value type: <class 'jsonargparse._namespace.Namespace'>

That is my intuition about the meaning of the error, perhaps there is something else going on.

Questions:
i) Is this the correct approach to load parameters from a checkpoint into a new model instantiation? We tried different approaches that were more problematic than this one (error-wise)
ii) If so, how can I manipulate the parser so that it would evaluate literally the torch.Tensor? Something similar to the CheckPointLoader constructor or equivalent to the flag type used in the .add_argument function in the classical ArgumentParser ?

I tried to look for possible approaches within the attributes from the LightningArgumentParser and ArgumentParser classes, however it did not make much sense to me.

Thank you in advance for your time!

[mauvilsa]

LightningCLI is a class that makes it simple to implement a CLI. But there are multiple ways of using it, and it is not clear what exactly you are doing, i.e. a minimal reproduction example. Also, the idea is that people implement the module classes with type hints, such that they are usable independent of any CLI. Without the CLI how would the checkpoint be given to the model class? LightningCLI is just a wrapper that interprets the parameters and type hints, but they should make sense without the CLI.

Using python-specific yaml syntax such as !CheckpointLoader ... is not the intended way of using it. The intention is that people can learn how to provide the different options by calling the CLI with --help and --print_config. For example, if LightningCLI is used with default options that includes subcommands, then if you run cli.py fit --help then you would see the --ckpt_path. This is a parameter of the Trainer.fit method that can be used to load a checkpoint, and there is no need to implement this logic in the model class. If you don't use subcommands, then the loading of a checkpoint would need to be implemented by you.

[LysSanzMoreta]

@mauvilsa Thank you very much for your quick reply. At the moment, it is a little complicated for me to provide a reasonable reproducible example at the moment since the code is full of nested classes and I am still trying to wrap my head around it.
Indeed, the ckpt_path argument could be helpful, since it loads the model into a new instance and then continues training (I have successfully tested it). What we were trying to do is to load the learnt parameters and then resume training whilst freezing the training of some selected ones. Is there a tutorial or help on how to manipulate the model loaded via the ckpt_path in a similar manner?
Thanks!

[mauvilsa]

Is there a tutorial or help on how to manipulate the model loaded via the ckpt_path in a similar manner?

Not that I know of.

[LysSanzMoreta]

So not a single clue on where to start?