loading weights from another model

[hkmztrk]

Hello,

I have two models, let's say ModelA and ModelB. ModelB is an extended version of ModelA. Let's say I trained Model A and saved it, now I want the load pretrained weights from ModelA to ModelB.

Would simply using ModelB.load_from_checkpoint() work?

class ModelA(pl.LightningModule):
    def __init__(self ):
        super().__init__()
        
        self.Conv1 = nn.Conv1d()
        self.Conv2 = nn.Conv1d()
        self.Maxpool = nn.MaxPool1d()
        self.Linear1 = nn.Linear()

    def forward(self, batches: dict) -> dict:
        inp = batches['ids'] 
        
        x = self.Conv1(inp)
        x = F.relu(x)
        x = self.Maxpool(x)
        x = self.Conv2(x)
        x = F.relu(x)
        x = self.Maxpool(x)
        out = self.Linear1(x)

        return logits

class ModelB(pl.LightningModule):
    def __init__(self ):
        super().__init__()
        
        self.Conv1 = nn.Conv1d()
        self.Conv2 = nn.Conv1d()
        self.Maxpool = nn.MaxPool1d()
        self.Linear1 = nn.Linear()
        self.Linear2 = nn.Linear()

    def forward(self, batches: dict) -> dict:
        inp = batches['ids'] 
        
        x = self.Conv1(inp)
        x = F.relu(x)
        x = self.Maxpool(x)
        x = self.Conv2(x)
        x = F.relu(x)
        x = self.Maxpool(x)
        x = self.Linear1(x)
        out = self.Linear2(x)

        return logits

[Abhishek-Prajapat]

I guess it's way too late for a reply but that won't work. what you could use is

base_model_for_b = ModelA()
base_model_for_b.load_state_dict('...')

class ModelB(pl.LightningModule):
    def __init__(self ):
        super().__init__()
        
        self.base_model = base_model_for_b
        self.Linear2 = nn.Linear()

    def forward(self, batches: dict) -> dict:
        inp = batches['ids'] 
        
        x = self.base_model(x)
        out = self.Linear2(x)

        return logits