Issues using a non-Lightning checkpoint in Lightning

pgyurov: <p>I have a checkpoint that was trained with a standard Pytorch implementation. The model used was <code>DeepLabV3Plus</code> from the <code>segmentation_models_pytorch</code> library.</p>
<p>I am trying to load the checkpoint with Pytorch <em>Lightning</em> but I am running into a few issues.</p>
<p>First I was getting <code>KeyError</code>s for <code>pytorch-lightning_version</code>, <code>global_step</code> and <code>epoch</code>. I set these to dummy values.</p>
<p>Then, I was getting the “Unexpected / Missing keys in state_dict” issue which I solved by prefixing the keys in the <code>state_dict</code> with <code>model.</code> (and removing the old ones).</p>
<p>But now I am stuck. I get the error:</p>
<pre><code class="lang-auto">self.trainer.predict_loop.load_state_dict(state_dict["predict_loop"])
KeyError: 'predict_loop'
</code></pre>
<p>Ok, so I added the <code>predict_loop</code> key to the state_dict and set it to <code>None</code>.</p>
<p>Now I get:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for DeepLabV3Plus:
        Unexpected key(s) in state_dict: "predict_loop". 
</code></pre>
<p>I assume the value of <code>predict_loop</code> can’t be <code>None</code> but then again that would make this error misleading.</p>
<p>Any ideas on how to solve this?</p>

tom: <p>I’m not exactly sure what you are doing from your description, but maybe loading the parameters to the model could be a good way, so <code>foo.model.load_state_dict()</code> if <code>foo.model</code> is the original model.</p>
<p>Best regards</p>
<p>Thomas</p>

pgyurov: <p>I am not sure I can load the checkpoint like that since Lightning takes care of that for me (unless there’s a way to override it?)</p>
<p>Here’s some context (reduced/simplified as much as possible). I am using the CLI and all my config is driven from a <code>config.yaml</code> file.</p>
<p><strong>main.py</strong></p>
<pre><code class="lang-python">cli = LightningCLI(..., run=False)
preds = cli.trainer.predict(cli.model, cli.datamodule, return_predictions=True, ckpt_path='&lt;my_ckpt.pth.tar&gt;')
</code></pre>
<p><strong>Model</strong></p>
<pre><code class="lang-python">import lightning.pytorch as pl
import segmentation_models_pytorch as smp
import torch.nn as nn

class DeepLabV3Plus(pl.LightningModule):
    def __init__(self, num_classes, encoder_name, encoder_depth,
                 encoder_weights, activation, loss_fn):
        super().__init__()

        self.model = smp.DeepLabV3Plus(encoder_name=encoder_name,
                              encoder_depth=encoder_depth,
                              encoder_weights=encoder_weights,
                              classes=num_classes,
                              activation=activation)

        self.loss_fn = nn.CrossEntropyLoss()
        self.save_hyperparameters()

    def configure_optimizers(self) -&gt; Any:
        return super().configure_optimizers()

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

    def training_step(self, batch, batch_idx):
        img, gt = batch
        output = self.model(img)
        loss = self.loss_fn(output, gt)
        self.log('train_loss', loss, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        img, gt = batch
        output = self.model(img)
        loss = self.loss_fn(output, gt)
        self.log('val_loss', loss, on_epoch=True, prog_bar=True)

    def test_step(self, batch, batch_idx):
        img, gt = batch
        output = self.model(img)
        loss = self.loss_fn(output, gt)
        self.log('test_loss', loss, on_epoch=True, prog_bar=True)

</code></pre>
<p>In a separate script I attempted to fix the checkpoint file like this:</p>
<pre><code class="lang-python">from collections import OrderedDict
import torch

# Load the checkpoint
checkpoint_path = "checkpoints/my_checkpoint.pth.tar"
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))

# create new checkpoint
modified_checkpoint = OrderedDict()
modified_checkpoint['state_dict'] = OrderedDict()

# prefix keys in state_dict
state_dict = checkpoint['state_dict'].copy()
for k, v in state_dict.items():
    k = f'model.{k}'
    modified_checkpoint['state_dict'][k] = v

# add missing keys
modified_checkpoint['pytorch-lightning_version'] = '0.0.0'
modified_checkpoint['global_step'] = None
modified_checkpoint['epoch'] = None
modified_checkpoint['state_dict']['predict_loop'] = None

# save
modified_checkpoint_path = "my_checkpoint_edited.pth.tar"
torch.save(modified_checkpoint, modified_checkpoint_path)
</code></pre>

tom: <p>How about loading the checkpoint to the (inner) model and then saving the full lightning model to get a lightning checkpoint?</p>

Ken_Pierce: <p>Did you ever figure this out? I’ve got a really similar problem except my model I trained didn’t use lightning but the script I’m using to generate output from the model is expecting a ckpt from a lightning module. I’m going to brazenly steal your deeplab3plus code above to see if I can get it rerunning as a lightning module. Thanks if you have any comments. Just learning some lightning but I’m in the endless, learn what you need to keep going cycle.</p>

pgyurov: <p>As I remember, I couldn’t get it working for my DeepLabV3Plus model, but I recently did (almost) the same thing for a Faster RCNN checkpoint and it worked.</p>
<p>This is more or less my code:</p>
<pre><code class="lang-auto">from collections import OrderedDict

import torch

# Load the checkpoint
checkpoint_path = "broken.ckpt"
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))

# create new checkpoint
modified_checkpoint = OrderedDict()
modified_checkpoint['state_dict'] = {}

# prefix keys in state_dict
for k, v in checkpoint.items():
    k = f'model.{k}'
    modified_checkpoint['state_dict'][k] = v

# add missing keys
modified_checkpoint['pytorch-lightning_version'] = '0.0.0'
modified_checkpoint['global_step'] = None
modified_checkpoint['epoch'] = None

# save
modified_checkpoint_path = "fixed.ckpt"
torch.save(modified_checkpoint, modified_checkpoint_path)

</code></pre>