How to transfer single GPU resume to dataparallel mode multi-GPU in train stage with Pytorch

duo_ma: <p>When I train a speech enhancement model with single GPU on the pytorch framework, Training stage is not finish. I stoped it because the speed is a little slow, Now I want to use data-parallel mode multiply GPU to continue training from resume model, what I should do?</p>

seungjun: <p>You can refer to the example code in this reply.</p>
<aside class="quote quote-modified" data-post="9" data-topic="97881">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/90ced4/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/save-checkpoints-trained-on-multi-gpus-for-load-on-single-gpu/97881/9">Save checkpoints trained on multi GPUs for load on single GPU</a> <a class="badge-category__wrapper " href="/c/distributed/12"><span data-category-id="12" data-drop-close="true" class="badge-category "><span class="badge-category__name">distributed</span></span></a>
  </div>
  <blockquote>
    Hi <a class="mention" href="/u/janine">@Janine</a>, 
This is not related to PyTorch version but the DataParallel (also DistributedDataParallel) class wrapper of PyTorch nn class models. 
DataParallel encloses the original model as it member variable, self.module. 
In case you need both single-GPU and multi-GPU model training, you can change saving/loading behavior with if statements. 
For example, 
if isinstance(G, nn.DataParallel):
    torch.save(G.module.state_dict(), model_save_name)
else:
    torch.save(G.state_dict(), model_save_…
  </blockquote>
</aside>


duo_ma: <p>First, thanks for your reply. Because single GPU training model, its state dict is different from Mutilply GPU training model. so it can’t load state dict from single GPU resume model on the multiply GPU training model.</p>

seungjun: <p>Hi <a class="mention" href="/u/duo_ma">@duo_ma</a>, sorry, I should have explained in more detail.</p>
<p><code>DataParallel</code> or <code>DistributedDataParallel</code> models calls it <code>state_dict</code> as <code>self.module.xxx</code> while single-gpu models call their <code>state_dict</code> without <code>self.module</code> element.</p>
<p>What I would recommend in general is always saving <code>state_dict</code> without <code>.module</code>. and loading it to model itself (single-gpu) or model.module (multi-gpu) accordingly.</p>
<pre><code class="lang-python"># saving
if isinstance(model, (nn.DataParallel, nn.DistributedDataParallel)):
    torch.save(model.module.state_dict(), model_save_name)
else:
    torch.save(model.state_dict(), model_save_name)
</code></pre>
<pre><code class="lang-python"># loading
model = nn.DataParallel(model, **gpu_device_arg) # multi-gpu model
if isinstance(model, (nn.DataParallel, nn.DistributedDataParallel)):
    model.module.load_state_dict(state_dict)    # your model will be loaded to multi-gpu model.
else:
    model.load_state_dict(state_dict)
</code></pre>
<p>In a case when you already saved multi-gpu model parameters as <code>.module.xxx</code> and loading to a single-gpu model, then you should do:</p>
<pre><code class="lang-python">model = nn.DataParallel(model, **gpu_device_arg) # make it multi-gpu
model.load_state_dict(state_dict)

model = model.module  # make it single-gpu
</code></pre>