How to load DataParallel model which trained using multiple GPUs?

Haizhuolaojisite: <p>Hi, I trained a model using 2 GPUs, and I want to make inference using trained model.</p>
<p><strong>1. How to load this parallelised model on GPU? or multiple GPU?</strong></p>
<p><strong>2. How to load this parallelised model on CPU?</strong></p>
<p>I find document mentioned the way to save the DataParallel model by add the “module”, but actually I successfully save the model in this way:</p>
<p>torch.save(best_model.state_dict(), bestmodel_path)<br>
Then I use the method below to load on CPU and it works:</p>
<pre><code class="lang-auto">def load_trainedModel_cpu(model_path,device=torch.device('cpu')):
    # Load trained resnet50
    model = models.resnet50(pretrained=False)
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(num_ftrs, 256), nn.ReLU(inplace=True), nn.Dropout(0.2),
        nn.Linear(256, 3))
    # Inference with CPU: use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.
    # original saved file with DataParallel
    state_dict = torch.load(model_path,map_location=device)
    # create new OrderedDict that does not contain `module.`
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] #remove 'module'
        new_state_dict[name] = v

    model.load_state_dict(new_state_dict)
    return model
</code></pre>
<h3>
<a name="saving-torchnndataparallel-models-1" class="anchor" href="#saving-torchnndataparallel-models-1"></a>Saving <code>torch.nn.DataParallel</code> Models</h3>
<p><strong>Save:</strong></p>
<p>torch.save(model.module.state_dict(), PATH)</p>
<p><strong>Load:</strong></p>
<h1>
<a name="load-to-whatever-device-you-want-2" class="anchor" href="#load-to-whatever-device-you-want-2"></a>Load to whatever device you want</h1>
<p><code>torch.nn.DataParallel</code> is a model wrapper that enables parallel GPU utilization. To save a <code>DataParallel</code> model generically, save the <code>model.module.state_dict()</code>. This way, you have the flexibility to load the model any way you want to any device you want.</p>
<p>Thank you!!</p>

ptrblck: <aside class="quote no-group" data-username="Haizhuolaojisite" data-post="1" data-topic="146005">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/haizhuolaojisite/48/47138_2.png" class="avatar"> Haizhuolaojisite:</div>
<blockquote>
<p>I find document mentioned the way to save the DataParallel model by add the “module”, but actually I successfully save the model in this way:</p>
</blockquote>
</aside>
<p>Saving the <code>model.module.state_dict()</code> would avoid having to manipulate the <code>state_dict</code> keys, but both approaches should work.</p>
<aside class="quote no-group" data-username="Haizhuolaojisite" data-post="1" data-topic="146005">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/haizhuolaojisite/48/47138_2.png" class="avatar"> Haizhuolaojisite:</div>
<blockquote>
<ol>
<li>How to load this parallelised model on GPU? or multiple GPU?</li>
</ol>
</blockquote>
</aside>
<p>Restore the model and wrap it again into <code>DistributedDataParallel</code> afterwards.</p>
<aside class="quote no-group" data-username="Haizhuolaojisite" data-post="1" data-topic="146005">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/haizhuolaojisite/48/47138_2.png" class="avatar"> Haizhuolaojisite:</div>
<blockquote>
<ol start="2">
<li>How to load this parallelised model on CPU?</li>
</ol>
</blockquote>
</aside>
<p>Store the <code>state_dict</code> on the CPU or use <code>map_location='cpu'</code> to load the <code>state_dict</code> onto the CPU.</p>