Unable to load pretrained weight into custom model in Pytorch Lightning

username: JohnDuke259
cooked: <p>Hello folks,</p>
<p>I want to retrain a custom model with my data. I can load the pretrained weights (.pth file) into the model in Pytorch and it runs but I want more functionality and refactored the code into Pytorch Lightning. I am having trouble loading the pretrained weight into the Pytorch Lightning model. The Pytorch Lightning code works but I have limited data and don’t have enough data to train from scratch.</p>
<p>Pytorch Model:</p>
<pre><code>class BDRAR(nn.Module):
      def __init__(self, stat=1):
          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
</code></pre>
<p>My Lightning model looks like this:</p>
<pre><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)
</code></pre>
<p>My run:</p>
<pre><code>KeyError: 'state_dict 
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)
</code></pre>
<p>I get the following error:</p>
<pre><code>    "keys = model.load_state_dict(checkpoint["state_dict"], strict=strict) **KeyError: 'state_dict**'"
</code></pre>
<p>I will appreciate any help. Thank you so much.</p>
<p><a class="mention" href="/u/williamfalcon">@williamFalcon</a></p>

username: ptrblck
cooked: <p>Would it work, if you load the <code>state_dict</code> into the internal <code>model</code> without using the Lightning method via <code>liteBDRARobject.model.load_state_dict()</code>?</p>

username: JohnDuke259
cooked: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a>, It works  when I load it to the pytorch model without using Lightning as shown below:</p>
<pre><code class="lang-auto">def main():
    net = BDRAR().cuda().train()
    print("Network Created")

    optimizer = optim.SGD([
        {'params': [param for name, param in net.named_parameters() if name[-4:] == 'bias'],
         'lr': 2 * args['lr']},
        {'params': [param for name, param in net.named_parameters() if name[-4:] != 'bias'],
         'lr': args['lr'], 'weight_decay': args['weight_decay']}
    ], momentum=args['momentum'])
    print("Training Started.....")
    if len(args['snapshot']) &gt; 0:
        print(('training resumes from \'%s\'' % args['snapshot']))
        net.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '.pth')),strict=False)
        optimizer.load_state_dict(torch.load(os.path.join(ckpt_path, exp_name, args['snapshot'] + '_optim.pth')))
        optimizer.param_groups[0]['lr'] = 2 * args['lr']
        optimizer.param_groups[1]['lr'] = args['lr']

    check_mkdir(ckpt_path)
    check_mkdir(os.path.join(ckpt_path, exp_name))
    open(log_path, 'w').write(str(args) + '\n\n')
    print("Net is: ")
    print(net)
</code></pre>

username: ptrblck
cooked: <p>I’m not familiar with Lightning enough, but are you able to access the “plain” PyTorch model from the Lightning wrapper and could then load the <code>state_dict</code> into it directly?</p>

username: JohnDuke259
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, many thanks for your reply. I appreciate it! I could access the plain PyTorch model from Lightning wrapper but I got an error when I tried to load the model.</p>
<p>I got help from another forum and I can load the pretrained model in Lightning now.</p>

username: erict
cooked: <p>Hi <a class="mention" href="/u/johnduke259">@JohnDuke259</a>,</p>
<p>can you please share your solution on how to load pytorch model into lightning module? I’am looking for the same thing.</p>
<p>Thanks</p>

username: JohnDuke259
cooked: <p>Hi <a class="mention" href="/u/erict">@erict</a>,</p>
<p>Sorry I just saw this notification.</p>
<p>In your Lightning Class:</p>
<pre><code class="lang-auto">class liteCalc(pl.LightningModule):

    def __init__(self):
        super(liteCalc, self).__init__()
        self.model = Calc()

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

    def forward(self, x):
        return self.model(x)
</code></pre>
<p>Outside Lightning class before you create your Trainer :</p>
<pre><code class="lang-auto">path = './ckpt/x.pth'
model = liteCalc()
model.load_model(path)
</code></pre>
<p>I hope this helps.</p>

username: erict
cooked: <p>Hi <a class="mention" href="/u/johnduke259">@JohnDuke259</a>,</p>
<p>yes this is working.<br>
Thanks</p>

username: JohnDuke259
cooked: <p>Hi <a class="mention" href="/u/erict">@erict</a> , you are welcome! <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:"></p>