Load weights from checkpoint using lightning-cli

[man-sean]

Hey 👋🏼

I trained a model using a regression loss and I want to use the learned weights as initialization for a GAN's generator (of the same type).
In regular training code I would call gan.generator.load_from_checkpoint(PATH) for that.
But I started using Lightning CLI and I'm a bit confused how to replicate this behavior.

My GAN model is structured like that:

class GAN(pl.LightningModule):
    def __init__(self,
                 generator: nn.Module,
                 critic: nn.Module,
                 [...]):
        super().__init__()
        self.save_hyperparameters()
        self.generator = generator
        [...]

And my config file has the following structure:

model:
  class_path: frameworks.GAN
  init_args:
    generator:
      class_path: models.Generator
      init_args:
        [...]
    critic:
      class_path: models.Critic
      init_args:
        [...]
    [...]

I would expect pl.LightningModule to have an argument checkpoint_path or something similar to perform load_from_checkpoint on its own in self.__init__. It seems it does not have such functionality.

How should I handle this staged training case using Lightning CLI?

Note: As I understand trainer.resume_from_checkpoint, it won't solve my problem, as it will not handle correctly the new structure of my model (and I don't want to resume the optimizers state etc.).

[mauvilsa]

I think you should add a parameter to your model to do this. All parameters you define get exposed in the CLI. Something like:

class MyModel(LightningModule):
    def __init__(
        self,
        weights_from_checkpoint_path: Optional[str] = None,
    ):
        if weights_from_checkpoint_path:
            checkpoint = torch.load(weights_from_checkpoint_path)
            self.load_state_dict(checkpoint['state_dict'])

[mshooter]

I get unexpected keys from the state_dict (e.g. self.model.encoder..)