How to provide a pretrained model as an argument to Lightning CLI?

[SirDavidLudwig]

I have a Lightning model that expects a pretrained encoder for one of its constructor parameters like so:

class MyModel(L.LightningModule):
    def __init__(self, encoder: nn.Module):
        super().__init__()
        self.encoder = encoder

cli = LightningCLI(MyModel, ...)

Using LightningCLI, how can I go about loading this pretrained encoder and passing it into the model via CLI arguments or a yaml config?

[mauvilsa]

By "pretrained" I suppose you mean that you want an encoder instantiated and some pretrained weights loaded, right? If so, the CLI by itself is not able to do that. The logic to load weights would need to be implemented in the module or the decoder and exposed as an init parameter.

[mauvilsa]

If you want to load a checkpoint, then you can use fit's ckpt_path which is exposed by default by the CLI.

[SirDavidLudwig]

Basically. I think the main idea that I had is that I would pass a path string to a pytorch model in the config somewhere. Is it possible to add some sort of transform to the argument parser to first load the model before running through the rest of Lightning CLI and validation?

For example, if I provided in the config:

model:
  encoder: "path/to/encoder.pt"

then could something happen somewhere to the argument parser that would perform the following transform:

args["model.encoder"] = torch.load(args["model.encoder"])

[mauvilsa]

The LightningCLI just interprets signature parameters, its type hints and defaults. Logic for loading of weights you should implement. So for example you could do:

class MyModel(L.LightningModule):
    def __init__(self, encoder: nn.Module, encoder_state_dict_path: Optional[str] = None):
        super().__init__()
        self.encoder = encoder
        if encoder_state_dict_path:
            encoder.load_state_dict(torch.load(encoder_state_dict_path))