How to Load Model from Checkpoint Without the Class for Ensemble

[aditya0by0]

Hello everyone,

I am looking for a way to load PyTorch Lightning models ( and/or even torch models) from checkpoints and perform a forward pass, particularly in cases where I don't have access to the model class implementation saved in the checkpoint. I want to make this process more generic.

Currently, I am trying to use load_from_checkpoint on LightningModule, but I’m encountering issues when the model class is not available. The challenge is that I only have access to the checkpoint files, and not the model class implementations.

Here’s an abstract example of my current implementation:

import torch
from lightning import LightningModule
from torch import nn
import torch.nn.functional as F

class _Base(LightningModule):
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = F.mse_loss(self.forward(x), y)
        self.log("train_loss", loss)
        return loss

class MyLightningModel1(_Base):
    def __init__(self, input_dim=50, out_dim=10):
        self.fc1 = nn.Linear(input_dim, out_dim)

    def forward(self, x):
        return self.fc1(x)

class MyLightningModel2(_Base):
    def __init__(self, input=50, out_layer=10):
        self.fc1 = nn.Linear(input, out_layer)

    def forward(self, x):
        return self.fc1(x)

class MyLightningModel3(_Base):
    """Some different model like graph with different input params"""

class EnsembleModel(_Base):
    def __init__(self, model_configs):
        self.models = {}
        
        self.fc1 = nn.Linear(len(model_configs), 10)

        for model_name, model_ckpt_path in model_configs.items():
            model = LightningModule.load_from_checkpoint(model_ckpt_path)
            model.eval()
            model.freeze()
            self.models[model_name] = model

    def forward(self, x):
        logits_list = []
        for model_name, model in self.models.items():
            logits_list.append(model(x))
        return self.fc1(torch.cat(logits_list, dim=1))

Problem & Question

In scenarios where I don’t have access to the class of the model saved in the checkpoint but still need to load it and perform a forward pass, how can I handle this? Is there a way to generalize the process so that the ensemble model can dynamically handle different types of models without explicitly referencing their classes?

Any insights or best practices for making the model loading and inference process more flexible would be greatly appreciated!

[aditya0by0]

@mauvilsa, Can you please help if you have any idea regarding the above.?

[mauvilsa]

From what I know, the model class implementation is not saved in the checkpoint. The checkpoint only stores the network weights and data to allow continuing a previous training. So not possible to do a forward if the model class is not available (installed) in the machine where the forward is run.

You should not be running LightningModule.load_from_checkpoint from inside a module.

To be able to easily switch between different models, I would recommend using the dependency injection pattern, see models-with-multiple-submodules. If you use this, the checkpoint will reference the submodules that were used during fit, and it should be possible to run load_from_checkpoint from the base class LightningModule without the need to specify more details. But all the used module class implementations must be available.

If the module class implementations are really not available, I guess you would need to use something like onnx. Thought, I have not used onnx, so can't advise more on that.