Unable to load custom pretrained weight in Pytorch Lightning
I want to retrain a custom model with my small dataset. I can load the pretrained weight (.pth) and run it in Pytorch. However, I need more functionalities and refactored the code to Pytorch lightning but I can't figure out how to load the pretrained weight into the Pytorch Lightning model.
Please see the details of my code below:
class BDRAR(nn.Module):
    def __init__(self):
        super(BDRAR, self).__init__()
        resnext = ResNeXt101()
        self.layer0 = resnext.layer0
        self.layer1 = resnext.layer1
        self.layer2 = resnext.layer2
        self.layer3 = resnext.layer3
        self.layer4 = resnext.layer4

Pytorch Lightning code:
class liteBDRAR(pl.LightningModule):
    def __init__(self):
        super(liteBDRAR, self).__init__()
        self.model = BDRAR()
        print('Model Created!')

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

Pytorch Lightning run:
    path = './ckpt/BDRAR/3000.pth'
    bdrar = liteBDRAR.load_from_checkpoint(path,  strict=False)
    trainer = pl.Trainer(fast_dev_run=True, gpus=1)
    trainer.fit(bdrar)


Error:
keys = model.load_state_dict(checkpoint[&quot;state_dict&quot;], strict=strict)
**KeyError: 'state_dict'**

I will appreciate any help.
Thank you.

It can be that your .pth file is already a state_dict. Try to load pretrained weight in your lightning class.
class liteBDRAR(pl.LightningModule):
    def __init__(self):
        super(liteBDRAR, self).__init__()
        self.model = BDRAR()
        print('Model Created!')

    def load_model(self, path):
        self.model.load_state_dict(torch.load(path, map_location='cuda:0'), strict=False)

path = './ckpt/BDRAR/3000.pth'
model = liteBDRAR()
model.load_model(path)


