Allow weight reuse in a different lightning module

[ziw-liu]

@mauvilsa I think we have a different issue than what #21246 fixes, where we want to load the state dict in a different lightning module. E.g. in the following we will have lr saved as a hyperparameter in the checkpoint from TrainingModule, but the InferenceModule does not take it as an argument:

import torch
from lightning.pytorch import LightningModule


class TrainingModule(LightningModule):
    def __init__(self, lr: float = 1e-3) -> None:
        super().__init__()
        self.model = torch.nn.Linear(16, 2)
        self.lr = lr

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer

    def training_step(
        self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int
    ) -> torch.Tensor:
        x, y = batch
        y_hat = self(x)
        loss = torch.nn.functional.cross_entropy(y_hat, y)
        self.log("train_loss", loss)
        return loss

    def train_dataloader(self) -> torch.utils.data.DataLoader:
        dataset = torch.utils.data.TensorDataset(
            torch.rand(32, 16), torch.randint(0, 2, (32,))
        )
        return torch.utils.data.DataLoader(dataset, batch_size=8)


class InferenceModule(LightningModule):
    def __init__(self) -> None:
        super().__init__()
        self.model = torch.nn.Linear(16, 2)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

    def predict_dataloader(self) -> torch.utils.data.DataLoader:
        dataset = torch.utils.data.TensorDataset(torch.rand(32, 16))
        return torch.utils.data.DataLoader(dataset, batch_size=128)

    def predict_step(
        self, batch: torch.Tensor, batch_idx: int, dataloader_idx: int = 0
    ) -> torch.Tensor:
        x = batch
        y_hat = self(x)
        return y_hat

Originally posted by @ziw-liu in #21116 (comment)

[mauvilsa]

The fix implemented in #21116 makes sure that when someone runs cli fit ... with module that saves hyperparameters, then when running a different subcommand, it is possible to give only the checkpoint path, e.g. cli test --ckpt_path ... without the need to specify again the correct parameters that the module requires, since these are found in the checkpoint. Naturally this requires the same module class to be used in fit and the other subcommand.

Option 1 - not good

To handle the case of different module classes one option could be to allow disabling the load of hyperparameters from the checkpoint. This can already be achieved by:

class MyCLI(LightningCLI):
    def _parse_ckpt_path(self):
        pass

This modifies a private method, which is not ideal. The simplest way to make this public would be a simple boolean flag in LightningCLI to disable this behavior.

However, I don't think this is a good proposal since it has the same flaw that #21116 attempted to fix. For example, assume that a module has a hidden_dim parameter and fit is run with a non-default hidden_dim. Running a different subcommand requires the CLI user to provide again the correct hidden_dim (and any other hyperparameter) even though they are already stored in the checkpoint.

Option 2 - not good

Instead of a boolean flag in LightningCLI loading of checkpoint hparams could be disabled with an additional argument, e.g. --load_chkp_hparams false. However, I think this is an even worse solution. First it deviates from arguments derived from the parameters of the Trainer methods. Second, it has the same flaw that #21116 attempted to fix.

Option 3 - actual proposal

Instead of disabling the loading of checkpoint hparams, we could add a way to customize this loading. People would need to subclass LightningCLI and implement a method for this customization. Something like:

class MyCLI(LightningCLI):
    def adapt_checkpoint_hparams(self, subcommand: str, chkp_hparams: dict) -> dict:
        ...

The subcommand is received in case the behavior should depend on the particular Trainer method called. If someone wants for some reason to completely disable load of checkpoint hparams, they can simply return {}. For the example described above, it would be:

chkp_hparams.pop("fit")
return chkp_hparams

If subclass module mode is enabled, then _class_path would need to be modified.