Does PyTorch Support Backward From DataParallel to DataParallel?

Eta_C: <p>For example:</p>
<pre><code class="lang-python">gpus = [0, 1]
model_1 = DataParallel(model_1, device_ids=gpus)  # such as resnet without classifier
model_2 = DataParallel(model_2, device_ids=gpus)  # a classifier contains dropout and linear 
opt.zero_grad()  # SGD optimizer
feat = model_1(x)
print("forward model_1")
prob = model_2(feat)
print("forward model_2")
loss = CrossEntropyLoss(prob, y)
print("forward loss")
loss.backward()
print("backward")
opt.step()
</code></pre>
<p>This program works fine on my toy dataset, which has only 200+ classes.<br>
But, when I put my large dataset, which has 2 million classes, it seems be stuck at backward step.<br>
Print</p>
<pre><code class="lang-python">forward model_1
forward model_2
forward loss
</code></pre>
<p>Addition:<br>
<code>nvidia-smi</code>:</p>
<ul>
<li>GPU-Util: 100%</li>
<li>Memory-Usage: about 15%</li>
</ul>
<p><code>top</code>:</p>
<ul>
<li>%CPU: 30</li>
<li>%MEM: 40%</li>
</ul>

Jamie_Donnelly: <p>I don’t see why you wouldn’t just wrap <code>model_1</code> and <code>model_2</code> in a single object, i.e.,</p>
<pre><code class="lang-auto">class full_model(nn.Module):
    def __init__():
        self.model1 = model1
        self.model2 = model2
        ...
    def forward(self,x):
        return self.model2(self.model1(x))
</code></pre>
<p>which can then be wrapped in a single <code>DataParallel</code> object. Depending on how you have instantiated your <code>opt</code> your current approach might not be working as intended, since your <code>opt</code> needs to be initialised with the <code>model</code> parameters. Which is why you should write it as follows:</p>
<pre><code class="lang-auto">model = full_model(...)
opt = torch.optim.Adam(model.parameters(),lr=lr)
model_dp = DataParallel(model, device_ids=gpus)
</code></pre>
<p>This way your parameters will be correctly updated when calling <code>opt.step()</code> since they likely weren’t before.</p>

Eta_C: <p><a class="mention" href="/u/jamie_donnelly">@Jamie_Donnelly</a> , Thanks for your reply.</p>
<aside class="quote no-group" data-username="Jamie_Donnelly" data-post="2" data-topic="175541">
<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/jamie_donnelly/48/56315_2.png" class="avatar"> Jamie_Donnelly:</div>
<blockquote>
<p>I don’t see why you wouldn’t just wrap <code>model_1</code> and <code>model_2</code> in a single object</p>
</blockquote>
</aside>
<p>I am building a modular project. When one use my project, there may be some other operators between <code>model_1</code> and <code>model_2</code>.</p>
<aside class="quote no-group" data-username="Jamie_Donnelly" data-post="2" data-topic="175541">
<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/jamie_donnelly/48/56315_2.png" class="avatar"> Jamie_Donnelly:</div>
<blockquote>
<p>Which is why you should write it as follows:</p>
<pre><code class="lang-auto">model = full_model(...)
opt = torch.optim.Adam(model.parameters(),lr=lr)
model_dp = DataParallel(model, device_ids=gpus)
</code></pre>
</blockquote>
</aside>
<p>Do you mean <code>DataParallel</code> must be after <code>optim</code>? I will try it, thank you again.<img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Jamie_Donnelly: <p>I think it will still work if you initialise an optimizer after a DataParellel model, but personally I think it’s better practice to initialise the optimizer on the base model itself, i.e., using <code>model.parameters()</code> in the <code>optim</code> constructor rather than <code>DataParallel(model).parameters()</code>.</p>
<p>Just make sure that your <code>optim</code> is initialised with the full model parameters, hence wrapping <code>model_1</code> and <code>model_2</code> into a single model.</p>