Very bad performance after loading from pth file

username: mariosconsta
cooked: <p>I have this function to save the model</p>
<pre><code class="lang-auto">save_checkpoint(
    {
        "epoch": epoch + 1,
        "state_dict": self.model.module.state_dict(),  # Model layers / weights
        "best_score": self.args[
            "best_pred"
        ],  # Save best score (scalar)
        "optimizer": self.optimizer.state_dict(),
    },
    is_best,  # Boolean. True if the predicted MAE is better than the last saved MAE
    self.args["save_path"],  # Where to save checkpoint
)


def save_checkpoint(state, is_best, save_path, filename="checkpoint.pth"):
    torch.save(state, "./" + str(save_path) + "/" + filename)
    if is_best:
        shutil.copyfile(
            "./" + str(save_path) + "/" + filename,
            "./" + str(save_path) + "/" + "model_best.pth",
        )
</code></pre>
<p>And this is how I load it:</p>
<pre><code class="lang-auto">def load_checkpoint(self, weight_path):
    if self.gpu_id == 0:
        print("### Loading Checkpoint ###")

    location = f"cuda:{self.gpu_id}"
    checkpoint = torch.load(weight_path, map_location=location)
    self.model.load_state_dict(checkpoint["state_dict"], strict=False)
</code></pre>
<p>When using <strong>strict=False</strong> I get no errors but the performance of the model is like it was never trained. When I use <strong>strick=True</strong> I get a <strong>alot</strong> of missing keys, e.g. Missing key(s) in state_dict: …</p>
<p>When I started working on this project I was using DP. Lately I migrated to DDP and it was working great. My last training session using DDP was successful and I wanted to run a validation run with it, so naturally I tried loading the pre-trained weights and all hell broke loose.</p>
<p>If anyone wants to see the actual error, its this:</p>
<pre><code class="lang-auto">Missing key(s) in state_dict: "module.conv1.weight", "module.in1.weight", "module.in1.bias", "module.conv2.weight", "module.in2.weight", "module.in2.bias", "module.layer1.0.conv1.weight", "module.layer1.0.in1.weight", "module.layer1.0.in1.bias", "module.layer1.0.conv2.weight", "module.layer1.0.in2.weight", "module.layer1.0.in2.bias", "module.layer1.0.conv3.weight", "module.layer1.0.in3.weight", "module.layer1.0.in3.bias", "module.layer1.0.downsample.0.weight", "module.layer1.0.downsample.1.weight", "module.layer1.0.downsample.1.bias", "module.layer1.1.conv1.weight", "module.layer1.1.in1.weight", "module.layer1.1.in1.bias", "module.layer1.1.conv2.weight", "module.layer1.1.in2.weight", "module.layer1.1.in2.bias", "module.layer1.1.conv3.weight", "module.layer1.1.in3.weight", "module.layer1.1.in3.bias", "module.layer1.2.conv1.weight", "module.layer1.2.in1.weight", "module.layer1.2.in1.bias", "module.layer1.2.conv2.weight", "module.layer1.2.in2.weight", "module.layer1.2.in2.bias", "module.layer1.2.conv3.weight", "module.layer1.2.in3.weight", "module.layer1.2.in3.bias", "module.layer1.3.conv1.weight", "module.layer1.3.in1.weight", "module.layer1.3.in1.bias", "module.layer1.3.conv2.weight", "module.layer1.3.in2.weight", "module.layer1.3.in2.bias", "module.layer1.3.conv3.weight", "module.layer1.3.in3.weight", "module.layer1.3.in3.bias", "module.transition1.0.0.weight", "module.transition1.0.1.weight", "module.transition1.0.1.bias", "module.transition1.1.0.0.weight", "module.transition1.1.0.1.weight", "module.transition1.1.0.1.bias", "module.stage2.0.branches.0.0.conv1.weight", "module.stage2.0.branches.0.0.in1.weight", "module.stage2.0.branches.0.0.in1.bias", "module.stage2.0.branches.0.0.conv2.weight", "module.stage2.0.branches.0.0.in2.weight", "module.stage2.0.branches.0.0.in2.bias", "module.stage2.0.branches.0.1.conv1.weight", "module.stage2.0.branches.0.1.in1.weight", "module.stage2.0.branches.0.1.in1.bias", "module.stage2.0.branches.0.1.conv2.weight", "module.stage2.0.branches.0.1.in2.weight", "module.stage2.0.branches.0.1.in2.bias", "module.stage2.0.branches.0.2.conv1.weight", "module.stage2.0.branches.0.2.in1.weight", "module.stage2.0.branches.0.2.in1.bias", "module.stage2.0.branches.0.2.conv2.weight", "module.stage2.0.branches.0.2.in2.weight", "module.stage2.0.branches.0.2.in2.bias", "module.stage2.0.branches.0.3.conv1.weight", "module.stage2.0.branches.0.3.in1.weight", "module.stage2.0.branches.0.3.in1.bias", "module.stage2.0.branches.0.3.conv2.weight", "module.stage2.0.branches.0.3.in2.weight", "module.stage2.0.branches.0.3.in2.bias", "module.stage2.0.branches.1.0.conv1.weight", "module.stage2.0.branches.1.0.in1.weight", "module.stage2.0.branches.1.0.in1.bias", "module.stage2.0.branches.1.0.conv2.weight", "module.stage2.0.branches.1.0.in2.weight", "module.stage2.0.branches.1.0.in2.bias", "module.stage2.0.branches.1.1.conv1.weight", "module.stage2.0.branches.1.1.in1.weight", "module.stage2.0.branches.1.1.in1.bias", "module.stage2.0.branches.1.1.conv2.weight", "module.stage2.0.branches.1.1.in2.weight", "module.stage2.0.branches.1.1.in2.bias", "module.stage2.0.branches.1.2.conv1.weight", "module.stage2.0.branches.1.2.in1.weight", "module.stage2.0.branches.1.2.in1.bias", "module.stage2.0.branches.1.2.conv2.weight", "module.stage2.0.branches.1.2.in2.weight", "module.stage2.0.branches.1.2.in2.bias", "module.stage2.0.branches.1.3.conv1.weight", "module.stage2.0.branches.1.3.in1.weight", "module.stage2.0.branches.1.3.in1.bias", "module.stage2.0.branches.1.3.conv2.weight", "module.stage2.0.branches.1.3.in2.weight", "module.stage2.0.branches.1.3.in2.bias", "module.stage2.0.fuse_layers.0.1.0.weight", "module.stage2.0.fuse_layers.0.1.1.weight", "module.stage2.0.fuse_layers.0.1.1.bias", "module.stage2.0.fuse_layers.1.0.0.0.weight", "module.stage2.0.fuse_layers.1.0.0.1.weight", "module.stage2.0.fuse_layers.1.0.0.1.bias", "module.transition2.2.0.0.weight", "module.transition2.2.0.1.weight", "module.transition2.2.0.1.bias", "module.stage3.0.branches.0.0.conv1.weight", "module.stage3.0.branches.0.0.in1.weight", "module.stage3.0.branches.0.0.in1.bias", "module.stage3.0.branches.0.0.conv2.weight", "module.stage3.0.branches.0.0.in2.weight", "module.stage3.0.branches.0.0.in2.bias", "module.stage3.0.branches.0.1.conv1.weight", "module.stage3.0.branches.0.1.in1.weight", "module.stage3.0.branches.0.1.in1.bias", "module.stage3.0.branches.0.1.conv2.weight", "module.stage3.0.branches.0.1.in2.weight", "module.stage3.0.branches.0.1.in2.bias", "module.stage3.0.branches.0.2.conv1.weight", "module.stage3.0.branches.0.2.in1.weight", "module.stage3.0.branches.0.2.in1.bias", "module.stage3.0.branches.0.2.conv2.weight", "module.stage3.0.branches.0.2.in2.weight", "module.stage3.0.branches.0.2.in2.bias", "module.stage3.0.branches.0.3.conv1.weight", "module.stage3.0.branches.0.3.in1.weight", "module.stage3.0.branches.0.3.in1.bias", "module.stage3.0.branches.0.3.conv2.weight", "module.stage3.0.branches.0.3.in2.weight", "module.stage3.0.branches.0.3.in2.bias", "module.stage3.0.branches.1.0.conv1.weight", "module.stage3.0.branches.1.0.in1.weight", "module.stage3.0.branches.1.0.in1.bias", "module.stage3.0.branches.1.0.conv2.weight", "module.stage3.0.branches.1.0.in2.weight", "module.stage3.0.branches.1.0.in2.bias", "module.stage3.0.branches.1.1.conv1.weight", "module.stage3.0.branches.1.1.in1.weight", "module.stage3.0.branches.1.1.in1.bias", "module.stage3.0.branches.1.1.conv2.weight", "module.stage3.0.branches.1.1.in2.weight", "module.stage3.0.branches.1.1.in2.bias", "module.stage3.0.branches.1.2.conv1.weight", "module.stage3.0.branches.1.2.in1.weight", "module.stage3.0.branches.1.2.in1.bias", "module.stage3.0.branches.1.2.conv2.weight", "module.stage3.0.branches.1.2.in2.weight", "module.stage3.0.branches.1.2.in2.bias", "module.stage3.0.branches.1.3.conv1.weight", "module.stage3.0.branches.1.3.in1.weight", "module.stage3.0.branches.1.3.in1.bias", "module.stage3.0.branches.1.3.conv2.weight", "module.stage3.0.branches.1.3.in2.weight", "module.stage3.0.branches.1.3.in2.bias", "module.stage3.0.branches.2.0.conv1.weight", "module.stage3.0.branches.2.0.in1.weight", "module.stage3.0.branches.2.0.in1.bias", "module.stage3.0.branches.2.0.conv2.weight", "module.stage3.0.branches.2.0.in2.weight", "module.stage3.0.branches.2.0.in2.bias", "module.stage3.0.branches.2.1.conv1.weight", "module.stage3.0.branches.2.1.in1.weight", "module.stage3.0.branches.2.1.in1.bias", "module.stage3.0.branches.2.1.conv2.weight", "module.stage3.0.branches.2.1.in2.weight", "module.stage3.0.branches.2.1.in2.bias", "module.stage3.0.branches.2.2.conv1.weight", "module.stage3.0.branches.2.2.in1.weight", "module.stage3.0.branches.2.2.in1.bias", "module.stage3.0.branches.2.2.conv2.weight", "module.stage3.0.branches.2.2.in2.weight", "module.stage3.0.branches.2.2.in2.bias", "module.stage3.0.branches.2.3.conv1.weight", "module.stage3.0.branches.2.3.in1.weight", "module.stage3.0.branches.2.3.in1.bias", "module.stage3.0.branches.2.3.conv2.weight", "module.stage3.0.branches.2.3.in2.weight", "module.stage3.0.branches.2.3.in2.bias", "module.stage3.0.fuse_layers.0.1.0.weight", "module.stage3.0.fuse_layers.0.1.1.weight", "module.stage3.0.fuse_layers.0.1.1.bias", "module.stage3.0.fuse_layers.0.2.0.weight", "module.stage3.0.fuse_layers.0.2.1.weight", "module.stage3.0.fuse_layers.0.2.1.bias", "module.stage3.0.fuse_layers.1.0.0.0.weight", "module.stage3.0.fuse_layers.1.0.0.1.weight", "module.stage3.0.fuse_layers.1.0.0.1.bias", "module.stage3.0.fuse_layers.1.2.0.weight", "module.stage3.0.fuse_layers.1.2.1.weight", "module.stage3.0.fuse_layers.1.2.1.bias", "module.stage3.0.fuse_layers.2.0.0.0.weight", "module.stage3.0.fuse_layers.2.0.0.1.weight", "module.stage3.0.fuse_layers.2.0.0.1.bias", "module.stage3.0.fuse_layers.2.0.1.0.weight", "module.stage3.0.fuse_layers.2.0.1.1.weight", "module.stage3.0.fuse_layers.2.0.1.1.bias", "module.stage3.0.fuse_layers.2.1.0.0.weight", "module.stage3.0.fuse_layers.2.1.0.1.weight", "module.stage3.0.fuse_layers.2.1.0.1.bias", "module.stage3.1.branches.0.0.conv1.weight", "module.stage3.1.branches.0.0.in1.weight", "module.stage3.1.branches.0.0.in1.bias", "module.stage3.1.branches.0.0.conv2.weight", "module.stage3.1.branches.0.0.in2.weight", "module.stage3.1.branches.0.0.in2.bias", "module.stage3.1.branches.0.1.conv1.weight", "module.stage3.1.branches.0.1.in1.weight", "module.stage3.1.branches.0.1.in1.bias", "module.stage3.1.branches.0.1.conv2.weight", "module.stage3.1.branches.0.1.in2.weight", "module.stage3.1.branches.0.1.in2.bias", "module.stage3.1.branches.0.2.conv1.weight", "module.stage3.1.branches.0.2.in1.weight", "module.stage3.1.branches.0.2.in1.bias", "module.stage3.1.branches.0.2.conv2.weight", "module.stage3.1.branches.0.2.in2.weight", "module.stage3.1.branches.0.2.in2.bias", "module.stage3.1.branches.0.3.conv1.weight", "module.stage3.1.branches.0.3.in1.weight", "module.stage3.1.branches.0.3.in1.bias", "module.stage3.1.branches.0.3.conv2.weight", "module.stage3.1.branches.0.3.in2.weight", "module.stage3.1.branches.0.3.in2.bias", "module.stage3.1.branches.1.0.conv1.weight", "module.stage3.1.branches.1.0.in1.weight", "module.stage3.1.branches.1.0.in1.bias", "module.stage3.1.branches.1.0.conv2.weight", "module.stage3.1.branches.1.0.in2.weight", "module.stage3.1.branches.1.0.in2.bias", "module.stage3.1.branches.1.1.conv1.weight", "module.stage3.1.branches.1.1.in1.weight", "module.stage3.1.branches.1.1.in1.bias", "module.stage3.1.branches.1.1.conv2.weight", "module.stage3.1.branches.1.1.in2.weight", "module.stage3.1.branches.1.1.in2.bias", "module.stage3.1.branches.1.2.conv1.weight", "module.stage3.1.branches.1.2.in1.weight", "module.stage3.1.branches.1.2.in1.bias", "module.stage3.1.branches.1.2.conv2.weight", "module.stage3.1.branches.1.2.in2.weight", "module.stage3.1.branches.1.2.in2.bias", "module.stage3.1.branches.1.3.conv1.weight", "module.stage3.1.branches.1.3.in1.weight", "module.stage3.1.branches.1.3.in1.bias", "module.stage3.1.branches.1.3.conv2.weight", "module.stage3.1.branches.1.3.in2.weight", "module.stage3.1.branches.1.3.in2.bias", "module.stage3.1.branches.2.0.conv1.weight", "module.stage3.1.branches.2.0.in1.weight", "module.stage3.1.branches.2.0.in1.bias", "module.stage3.1.branches.2.0.conv2.weight", "module.stage3.1.branches.2.0.in2.weight", "module.stage3.1.branches.2.0.in2.bias", "module.stage3.1.branches.2.1.conv1.weight", "module.stage3.1.branches.2.1.in1.weight", "module.stage3.1.branches.2.1.in1.bias", "module.stage3.1.branches.2.1.conv2.weight", "module.stage3.1.branches.2.1.in2.weight", "module.stage3.1.branches.2.1.in2.bias", "module.stage3.1.branches.2.2.conv1.weight", "module.stage3.1.branches.2.2.in1.weight", "module.stage3.1.branches.2.2.in1.bias", "module.stage3.1.branches.2.2.conv2.weight", "module.stage3.1.branches.2.2.in2.weight", "module.stage3.1.branches.2.2.in2.bias", "module.stage3.1.branches.2.3.conv1.weight", "module.stage3.1.branches.2.3.in1.weight", "module.stage3.1.branches.2.3.in1.bias", "module.stage3.1.branches.2.3.conv2.weight", "module.stage3.1.branches.2.3.in2.weight", "module.stage3.1.branches.2.3.in2.bias", "module.stage3.1.fuse_layers.0.1.0.weight", "module.stage3.1.fuse_layers.0.1.1.weight", "module.stage3.1.fuse_layers.0.1.1.bias", "module.stage3.1.fuse_layers.0.2.0.weight", "module.stage3.1.fuse_layers.0.2.1.weight", "module.stage3.1.fuse_layers.0.2.1.bias", "module.stage3.1.fuse_layers.1.0.0.0.weight", "module.stage3.1.fuse_layers.1.0.0.1.weight", "module.stage3.1.fuse_layers.1.0.0.1.bias", "module.stage3.1.fuse_layers.1.2.0.weight", "module.stage3.1.fuse_layers.1.2.1.weight", "module.stage3.1.fuse_layers.1.2.1.bias", "module.stage3.1.fuse_layers.2.0.0.0.weight", "module.stage3.1.fuse_layers.2.0.0.1.weight", "module.stage3.1.fuse_layers.2.0.0.1.bias", "module.stage3.1.fuse_layers.2.0.1.0.weight", "module.stage3.1.fuse_layers.2.0.1.1.weight", "module.stage3.1.fuse_layers.2.0.1.1.bias", "module.stage3.1.fuse_layers.2.1.0.0.weight", "module.stage3.1.fuse_layers.2.1.0.1.weight", "module.stage3.1.fuse_layers.2.1.0.1.bias", "module.stage3.2.branches.0.0.conv1.weight", "module.stage3.2.branches.0.0.in1.weight", "module.stage3.2.branches.0.0.in1.bias", "module.stage3.2.branches.0.0.conv2.weight", "module.stage3.2.branches.0.0.in2.weight", "module.stage3.2.branches.0.0.in2.bias", "module.stage3.2.branches.0.1.conv1.weight", "module.stage3.2.branches.0.1.in1.weight", "module.stage3.2.branches.0.1.in1.bias", "module.stage3.2.branches.0.1.conv2.weight", "module.stage3.2.branches.0.1.in2.weight", "module.stage3.2.branches.0.1.in2.bias", "module.stage3.2.branches.0.2.conv1.weight", "module.stage3.2.branches.0.2.in1.weight", "module.stage3.2.branches.0.2.in1.bias", "module.stage3.2.branches.0.2.conv2.weight", "module.stage3.2.branches.0.2.in2.weight", "module.stage3.2.branches.0.2.in2.bias", "module.stage3.2.branches.0.3.conv1.weight", "module.stage3.2.branches.0.3.in1.weight", "module.stage3.2.branches.0.3.in1.bias", "module.stage3.2.branches.0.3.conv2.weight", "module.stage3.2.branches.0.3.in2.weight", "module.stage3.2.branches.0.3.in2.bias", "module.stage3.2.branches.1.0.conv1.weight", "module.stage3.2.branches.1.0.in1.weight", "module.stage3.2.branches.1.0.in1.bias", "module.stage3.2.branches.1.0.conv2.weight", "module.stage3.2.branches.1.0.in2.weight", "module.stage3.2.branches.1.0.in2.bias", "module.stage3.2.branches.1.1.conv1.weight", "module.stage3.2.branches.1.1.in1.weight", "module.stage3.2.branches.1.1.in1.bias", "module.stage3.2.branches.1.1.conv2.weight", "module.stage3.2.branches.1.1.in2.weight", "module.stage3.2.branches.1.1.in2.bias", "module.stage3.2.branches.1.2.conv1.weight", "module.stage3.2.branches.1.2.in1.weight", "module.stage3.2.branches.1.2.in1.bias", "module.stage3.2.branches.1.2.conv2.weight", "module.stage3.2.branches.1.2.in2.weight", "module.stage3.2.branches.1.2.in2.bias", "module.stage3.2.branches.1.3.conv1.weight", "module.stage3.2.branches.1.3.in1.weight", "module.stage3.2.branches.1.3.in1.bias", "module.stage3.2.branches.1.3.conv2.weight", "module.stage3.2.branches.1.3.in2.weight", "module.stage3.2.branches.1.3.in2.bias", "module.stage3.2.branches.2.0.conv1.weight", "module.stage3.2.branches.2.0.in1.weight", "module.stage3.2.branches.2.0.in1.bias", "module.stage3.2.branches.2.0.conv2.weight", "module.stage3.2.branches.2.0.in2.weight", "module.stage3.2.branches.2.0.in2.bias", "module.stage3.2.branches.2.1.conv1.weight", "module.stage3.2.branches.2.1.in1.weight", "module.stage3.2.branches.2.1.in1.bias", "module.stage3.2.branches.2.1.conv2.weight", "module.stage3.2.branches.2.1.in2.weight", "module.stage3.2.branches.2.1.in2.bias", "module.stage3.2.branches.2.2.conv1.weight", "module.stage3.2.branches.2.2.in1.weight", "module.stage3.2.branches.2.2.in1.bias", "module.stage3.2.branches.2.2.conv2.weight", "module.stage3.2.branches.2.2.in2.weight", "module.stage3.2.branches.2.2.in2.bias", "module.stage3.2.branches.2.3.conv1.weight", "module.stage3.2.branches.2.3.in1.weight", "module.stage3.2.branches.2.3.in1.bias", "module.stage3.2.branches.2.3.conv2.weight", "module.stage3.2.branches.2.3.in2.weight", "module.stage3.2.branches.2.3.in2.bias", "module.stage3.2.fuse_layers.0.1.0.weight", "module.stage3.2.fuse_layers.0.1.1.weight", "module.stage3.2.fuse_layers.0.1.1.bias", "module.stage3.2.fuse_layers.0.2.0.weight", "module.stage3.2.fuse_layers.0.2.1.weight", "module.stage3.2.fuse_layers.0.2.1.bias", "module.stage3.2.fuse_layers.1.0.0.0.weight", "module.stage3.2.fuse_layers.1.0.0.1.weight", "module.stage3.2.fuse_layers.1.0.0.1.bias", "module.stage3.2.fuse_layers.1.2.0.weight", "module.stage3.2.fuse_layers.1.2.1.weight", "module.stage3.2.fuse_layers.1.2.1.bias", "module.stage3.2.fuse_layers.2.0.0.0.weight", "module.stage3.2.fuse_layers.2.0.0.1.weight", "module.stage3.2.fuse_layers.2.0.0.1.bias", "module.stage3.2.fuse_layers.2.0.1.0.weight", "module.stage3.2.fuse_layers.2.0.1.1.weight", "module.stage3.2.fuse_layers.2.0.1.1.bias", "module.stage3.2.fuse_layers.2.1.0.0.weight", "module.stage3.2.fuse_layers.2.1.0.1.weight", "module.stage3.2.fuse_layers.2.1.0.1.bias", "module.stage3.3.branches.0.0.conv1.weight", "module.stage3.3.branches.0.0.in1.weight", "module.stage3.3.branches.0.0.in1.bias", "module.stage3.3.branches.0.0.conv2.weight", "module.stage3.3.branches.0.0.in2.weight", "module.stage3.3.branches.0.0.in2.bias", "module.stage3.3.branches.0.1.conv1.weight", "module.stage3.3.branches.0.1.in1.weight", "module.stage3.3.branches.0.1.in1.bias", "module.stage3.3.branches.0.1.conv2.weight", "module.stage3.3.branches.0.1.in2.weight", "module.stage3.3.branches.0.1.in2.bias", "module.stage3.3.branches.0.2.conv1.weight", "module.stage3.3.branches.0.2.in1.weight", "module.stage3.3.branches.0.2.in1.bias", "module.stage3.3.branches.0.2.conv2.weight", "module.stage3.3.branches.0.2.in2.weight", "module.stage3.3.branches.0.2.in2.bias", "module.stage3.3.branches.0.3.conv1.weight", "module.stage3.3.branches.0.3.in1.weight", "module.stage3.3.branches.0.3.in1.bias", "module.stage3.3.branches.0.3.conv2.weight", "module.stage3.3.branches.0.3.in2.weight", "module.stage3.3.branches.0.3.in2.bias", "module.stage3.3.branches.1.0.conv1.weight", "module.stage3.3.branches.1.0.in1.weight", "module.stage3.3.branches.1.0.in1.bias", "module.stage3.3.branches.1.0.conv2.weight", "module.stage3.3.branches.1.0.in2.weight", "module.stage3.3.branches.1.0.in2.bias", "module.stage3.3.branches.1.1.conv1.weight", "module.stage3.3.branches.1.1.in1.weight", "module.stage3.3.branches.1.1.in1.bias", "module.stage3.3.branches.1.1.conv2.weight", "module.stage3.3.branches.1.1.in2.weight", "module.stage3.3.branches.1.1.in2.bias", "module.stage3.3.branches.1.2.conv1.weight", "module.stage3.3.branches.1.2.in1.weight", "module.stage3.3.branches.1.2.in1.bias", "module.stage3.3.branches.1.2.conv2.weight", "module.stage3.3.branches.1.2.in2.weight", "module.stage3.3.branches.1.2.in2.bias", "module.stage3.3.branches.1.3.conv1.weight", "module.stage3.3.branches.1.3.in1.weight", "module.stage3.3.branches.1.3.in1.bias", "module.stage3.3.branches.1.3.conv2.weight", "module.stage3.3.branches.1.3.in2.weight", "module.stage3.3.branches.1.3.in2.bias", "module.stage3.3.branches.2.0.conv1.weight", "module.stage3.3.branches.2.0.in1.weight", "module.stage3.3.branches.2.0.in1.bias", "module.stage3.3.branches.2.0.conv2.weight", "module.stage3.3.branches.2.0.in2.weight", "module.stage3.3.branches.2.0.in2.bias", "module.stage3.3.branches.2.1.conv1.weight", "module.stage3.3.branches.2.1.in1.weight", "module.stage3.3.branches.2.1.in1.bias", "module.stage3.3.branches.2.1.conv2.weight", "module.stage3.3.branches.2.1.in2.weight", "module.stage3.3.branches.2.1.in2.bias", "module.stage3.3.branches.2.2.conv1.weight", "module.stage3.3.branches.2.2.in1.weight", "module.stage3.3.branches.2.2.in1.bias", "module.stage3.3.branches.2.2.conv2.weight", "module.stage3.3.branches.2.2.in2.weight", "module.stage3.3.branches.2.2.in2.bias", "module.stage3.3.branches.2.3.conv1.weight", "module.stage3.3.branches.2.3.in1.weight", "module.stage3.3.branches.2.3.in1.bias", "module.stage3.3.branches.2.3.conv2.weight", "module.stage3.3.branches.2.3.in2.weight", "module.stage3.3.branches.2.3.in2.bias", "module.stage3.3.fuse_layers.0.1.0.weight", "module.stage3.3.fuse_layers.0.1.1.weight", "module.stage3.3.fuse_layers.0.1.1.bias", "module.stage3.3.fuse_layers.0.2.0.weight", "module.stage3.3.fuse_layers.0.2.1.weight", "module.stage3.3.fuse_layers.0.2.1.bias", "module.stage3.3.fuse_layers.1.0.0.0.weight", "module.stage3.3.fuse_layers.1.0.0.1.weight", "module.stage3.3.fuse_layers.1.0.0.1.bias", "module.stage3.3.fuse_layers.1.2.0.weight", "module.stage3.3.fuse_layers.1.2.1.weight", "module.stage3.3.fuse_layers.1.2.1.bias", "module.stage3.3.fuse_layers.2.0.0.0.weight", "module.stage3.3.fuse_layers.2.0.0.1.weight", "module.stage3.3.fuse_layers.2.0.0.1.bias", "module.stage3.3.fuse_layers.2.0.1.0.weight", "module.stage3.3.fuse_layers.2.0.1.1.weight", "module.stage3.3.fuse_layers.2.0.1.1.bias", "module.stage3.3.fuse_layers.2.1.0.0.weight", "module.stage3.3.fuse_layers.2.1.0.1.weight", "module.stage3.3.fuse_layers.2.1.0.1.bias", "module.transition3.3.0.0.weight", "module.transition3.3.0.1.weight", "module.transition3.3.0.1.bias", 
</code></pre>
<p><strong>and this</strong></p>
<pre><code class="lang-auto">Unexpected key(s) in state_dict: "conv1.weight", "in1.weight", "in1.bias", "conv2.weight", "in2.weight", "in2.bias", "layer1.0.conv1.weight", "layer1.0.in1.weight", "layer1.0.in1.bias", 
</code></pre>
<p>This is not the full error, I cut some of it out due to character limit.</p>
<p>I hope I can solve this without re-training the model. It took almost a week to finish.</p>

username: mariosconsta
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0d29a3365abec1be13f0a378b3895e7b5960de4.png" data-download-href="https://discuss.pytorch.org/uploads/default/a0d29a3365abec1be13f0a378b3895e7b5960de4" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/0/a0d29a3365abec1be13f0a378b3895e7b5960de4.png" alt="image" data-base62-sha1="mWHBJUJzUdUYirwSGduYav8FAXO" width="495" height="500" data-dominant-color="2D3237"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">717×724 17.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
looking inside the state_dict key shows that the layers are there but they are missing the “module” part for some reason…</p>
<p>Doing this fixed it:</p>
<pre><code class="lang-auto">        for k, v in state_dict.items():
            if not k.startswith("module."):
                k = f"module.{k}"
            new_state_dict[k] = v
        checkpoint["state_dict"] = new_state_dict
        self.model.load_state_dict(checkpoint["state_dict"])
</code></pre>
<p>I basically added what it was looking for. I don’t know what happened initially but it works.</p>

username: ptrblck
cooked: <p>The additional <code>.module</code> keys were added by <code>nn.DataParallel</code> and manipulating the <code>state_dict</code> is a valid approach to fix it. Use <code>strict=False</code> only if you are sure you can ignore missing or unexpected keys.</p>

username: mariosconsta
cooked: <p>I must have messed something up since the .module was missing from the state_dict keys! Generally when using DDP, is this a valid basic approach to saving / loading models?</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L62">
  <header class="source">

      <a href="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L62" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L62" target="_blank" rel="noopener nofollow ugc">pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L62</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="52" style="counter-reset: li-counter 51 ;">
          <li></li>
          <li>def _run_epoch(self, epoch):</li>
          <li>    b_sz = len(next(iter(self.train_data))[0])</li>
          <li>    print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")</li>
          <li>    self.train_data.sampler.set_epoch(epoch)</li>
          <li>    for source, targets in self.train_data:</li>
          <li>        source = source.to(self.gpu_id)</li>
          <li>        targets = targets.to(self.gpu_id)</li>
          <li>        self._run_batch(source, targets)</li>
          <li></li>
          <li class="selected">def _save_snapshot(self, epoch):</li>
          <li>    snapshot = {</li>
          <li>        "MODEL_STATE": self.model.module.state_dict(),</li>
          <li>        "EPOCHS_RUN": epoch,</li>
          <li>    }</li>
          <li>    torch.save(snapshot, self.snapshot_path)</li>
          <li>    print(f"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}")</li>
          <li></li>
          <li>def train(self, max_epochs: int):</li>
          <li>    for epoch in range(self.epochs_run, max_epochs):</li>
          <li>        self._run_epoch(epoch)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39">
  <header class="source">

      <a href="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39" target="_blank" rel="noopener nofollow ugc">pytorch/examples/blob/main/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="29" style="counter-reset: li-counter 28 ;">
          <li>    self.optimizer = optimizer</li>
          <li>    self.save_every = save_every</li>
          <li>    self.epochs_run = 0</li>
          <li>    self.snapshot_path = snapshot_path</li>
          <li>    if os.path.exists(snapshot_path):</li>
          <li>        print("Loading snapshot")</li>
          <li>        self._load_snapshot(snapshot_path)</li>
          <li></li>
          <li>    self.model = DDP(self.model, device_ids=[self.gpu_id])</li>
          <li></li>
          <li class="selected">def _load_snapshot(self, snapshot_path):</li>
          <li>    loc = f"cuda:{self.gpu_id}"</li>
          <li>    snapshot = torch.load(snapshot_path, map_location=loc)</li>
          <li>    self.model.load_state_dict(snapshot["MODEL_STATE"])</li>
          <li>    self.epochs_run = snapshot["EPOCHS_RUN"]</li>
          <li>    print(f"Resuming training from snapshot at Epoch {self.epochs_run}")</li>
          <li></li>
          <li>def _run_batch(self, source, targets):</li>
          <li>    self.optimizer.zero_grad()</li>
          <li>    output = self.model(source)</li>
          <li>    loss = F.cross_entropy(output, targets)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: ptrblck
cooked: <p>Yes, storing the <code>state_dict</code> from the internal <code>.module</code> allows you to directly load it into a “standard” (i.e. non-DDP) model and is also described <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-torch-nn-dataparallel-models">here</a>. Otherwise, if you store the <code>state_dict</code> directly from the <code>DDP</code> wrapper you would either need to remove the <code>.module</code> keys before loading the <code>state_dict</code> into a standard model or you could load it into a <code>DDP</code> model.</p>