Loss becomes higher after resuming from DDP model checkpoint

111248: <h1><a name="problem-1" class="anchor" href="#problem-1"></a>Problem</h1>
<p>Recently, I have tried to use DDP. I just followed the tutorial to change my original code for DP.<br>
I accidentally stopped my DDP training, so I planned to resume model as what I did before. However, loss becomes much higher after resuming, I think it might be some error in my code.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/0/d05f13fe78adc5f895c4d02cd3e22ba073165b57.png" data-download-href="https://discuss.pytorch.org/uploads/default/d05f13fe78adc5f895c4d02cd3e22ba073165b57" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d05f13fe78adc5f895c4d02cd3e22ba073165b57_2_690x190.png" alt="image" data-base62-sha1="tJl16dyeYFTnBIVxYHT2G12Cu6r" width="690" height="190" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d05f13fe78adc5f895c4d02cd3e22ba073165b57_2_690x190.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d05f13fe78adc5f895c4d02cd3e22ba073165b57_2_1035x285.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d05f13fe78adc5f895c4d02cd3e22ba073165b57_2_1380x380.png 2x" data-dominant-color="FCFCFC"><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">2167×598 36.4 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>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/f/4f9f1e9322b6b8923a6153162dcdb8a16183cc44.png" data-download-href="https://discuss.pytorch.org/uploads/default/4f9f1e9322b6b8923a6153162dcdb8a16183cc44" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f9f1e9322b6b8923a6153162dcdb8a16183cc44_2_690x175.png" alt="image" data-base62-sha1="bmmDbjX0uOkMGOxUFF3pc6oEdW4" width="690" height="175" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f9f1e9322b6b8923a6153162dcdb8a16183cc44_2_690x175.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f9f1e9322b6b8923a6153162dcdb8a16183cc44_2_1035x262.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f9f1e9322b6b8923a6153162dcdb8a16183cc44_2_1380x350.png 2x" data-dominant-color="FDFCFC"><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">2447×622 40.1 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></p>
<hr>
<h1><a name="here-are-my-codes-2" class="anchor" href="#here-are-my-codes-2"></a>Here are my codes</h1>
<h2><a name="save-model-3" class="anchor" href="#save-model-3"></a>save model</h2>
<pre><code class="lang-auto"># code for resuming after validation
if args.local_rank == 0:
    tf_writer.add_scalar('acc/test_top1_best', best_prec1, epoch)

    output_best = 'Best Prec@1: %.3f\n' % (best_prec1)
    print(output_best)
    log_training.write(output_best + '\n')
    log_training.flush()

    save_checkpoint({
        'epoch': epoch + 1,
        'arch': args.arch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'best_prec1': best_prec1,
    }, is_best)
</code></pre>
<pre><code class="lang-auto">def save_checkpoint(state, is_best):
    filename = '%s/%s/ckpt.pth.tar' % (args.root_model, args.store_name)
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, filename.replace('pth.tar', 'best.pth.tar'))
</code></pre>
<p>I don’t save scheduler because I manually set the learning rate for the optimizer. And I will set learning rate before training.</p>
<pre><code class="lang-auto">adjust_learning_rate(optimizer, epoch, args.lr_type, args.lr_steps)
</code></pre>
<h2><a name="resume-model-4" class="anchor" href="#resume-model-4"></a>resume model</h2>
<pre><code class="lang-auto">if args.resume:
    if os.path.isfile(args.resume):
        checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
        pretrained_dict = checkpoint['state_dict']
        new_state_dict = OrderedDict()
        for k, v in pretrained_dict.items():
            if '.total' not in k:
                name = k[7:]  # remove 'module.'
                # name = name.replace('.net', '')
                new_state_dict[name] = v
        model.load_state_dict(new_state_dict)

        model=torch.nn.parallel.DistributedDataParallel(model,
                                                device_ids=[local_rank],
                                                output_device=local_rank)

        if 'epoch' in checkpoint.keys():
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            
            # get_optim_policies is a fuction to set optimizer policies
            optimizer = torch.optim.SGD(get_optim_policies(model), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

            optimizer.load_state_dict(checkpoint['optimizer'])
            print(("=&gt; loaded checkpoint '{}' (epoch {})"
                .format(args.evaluate, checkpoint['epoch'])))
            print(("=&gt; best top1 '{}'".format(best_prec1)))
</code></pre>
<p>The above codes work well when I use DP. Are there anything I miss while using DDP?</p>

mrshenli: <p>When you call <code>save_checkpoint</code>, is the <code>model</code> var a DDP instance? If yes, you might need to save <code>model.module</code> instead? But I don’t think that’s the reason for the error jump. When you load the module, if you do not use DDP or DP (just a local model), is the loss after recovery as expected?</p>
<p>I might miss sth, looks like in the “resume model” part, the model state is loaded to CPU and not moved to <code>local_rank</code> before passing to DDP ctors? Or is the model already on the correct device before loading state dict?</p>