Why does requires_grad = True lead to missing keys in the state_dict only when using DataParallel?

23vil: <p>Hi everyone! I have a problem understanding what happens to state_dicts, when using DataParallel. Once I apply DataParallel to a model with <code>requires_grad = True</code>, then the <code>state_dict().keys()</code> are empty.<br>
My code has these two classes:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision
from torchvision import models

class LTE(torch.nn.Module):
    def __init__(self, requires_grad=True):
        super(LTE, self).__init__()
        vgg_pretrained_features = models.vgg19(pretrained=True).features
        self.slice1 = torch.nn.Sequential()

        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x]) 
        if not requires_grad:
            for param in self.slice1.parameters():
                param.requires_grad = requires_grad
</code></pre>
<pre><code class="lang-auto">class TTSR(nn.Module):    
    def __init__(self):
        super(TTSR, self).__init__()
        self.LTE      = LTE(requires_grad=True) 
        self.LTE_copy = LTE(requires_grad=False)
    def forward(self, sr=None):
        print("LTE "+str(torch.cuda.current_device())+"    "+str(self.LTE.state_dict().keys()))
        print("LTE_copy "+str(torch.cuda.current_device())+"    "+str(self.LTE_copy.state_dict().keys()))
</code></pre>
<p>If I now initialize a TTSR model and run the forward function:</p>
<pre><code class="lang-auto">device = torch.device('cuda')
_model = TTSR().to(device)
a = _model(sr=1)
</code></pre>
<p>Then, this is my output:</p>
<pre><code class="lang-auto">LTE 0    odict_keys(['slice1.0.weight', 'slice1.0.bias'])
LTE_copy 0    odict_keys(['slice1.0.weight', 'slice1.0.bias'])
</code></pre>
<p>This is fine so far! But if I use DataParallel now (with two GPUs),</p>
<pre><code class="lang-auto">model = nn.DataParallel(_model, list(range(2)))
a = model(sr=1)
</code></pre>
<p>my output is as follows:</p>
<pre><code class="lang-auto">LTE 0    odict_keys([ ])
LTE_copy 0    odict_keys(['slice1.0.weight', 'slice1.0.bias'])
LTE 1    odict_keys([ ])
LTE_copy 1    odict_keys([ ])
</code></pre>
<ul>
<li>Why does requires_grad = True lead to those Keys [‘slice1.0.weight’, ‘slice1.0.bias’] missing in the state_dict? This only happens once I send that model to DataParallel.</li>
<li>Did I understand it correct, that DataParallel always has one main GPU where the model is stored in (<a href="https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551" class="inline-onebox">DataParallel imbalanced memory usage</a>) and therefore the state_dicts of LTE1 &amp; LTE_copy1 are empty?</li>
</ul>

pbelevich: <p>I suspect that you don’t see the keys because of <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/replicate.py#L130" rel="noopener nofollow ugc">this line</a> which is called from <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L172" rel="noopener nofollow ugc">DataParallel.forward</a>. Basically you should not try to access replicas on different GPUs directly, because they are handled in different way than regular modules. Also, we recommend to use <a href="https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel" rel="noopener nofollow ugc"><code>DistributedDataParallel</code></a>, instead of this class, to do multi-GPU training, even if there is only a single node. See: <a href="https://pytorch.org/docs/stable/notes/cuda.html#cuda-nn-ddp-instead" rel="noopener nofollow ugc">Use nn.parallel.DistributedDataParallel instead of multiprocessing or nn.DataParallel</a> and <a href="https://pytorch.org/docs/stable/notes/ddp.html#ddp" rel="noopener nofollow ugc">Distributed Data Parallel</a>.</p>