Changing model structure in nn.DataParallel

username: NoobCoder
cooked: <p>Hello,</p>
<p>I want to change the model layers of my previously trained models. For example, I have a MobileNet trained on my data, and now I want to remove the FC layers. For this purpose, I replace the FC layer with a class called Identity, which you can find the definition below. (I took this suggestion from one of the admins in here, but I do not remember the post)</p>
<pre><code class="lang-auto">class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x
</code></pre>
<p>The procedures are this:</p>
<ol>
<li>Define a MobileNet by using <code>torchvision.model</code></li>
<li>Move the model to <code>nn.DataParallel</code></li>
<li>Load the previously trained model</li>
<li>Replace FC layers with <code>Identity</code> class</li>
</ol>
<p>Now, the problem here is that the model is not fully replaced, and I think it is because there are other copies on other GPUs as well.</p>
<p>Is there any solution to change all the models? This is the result of <code>print(model)</code>.</p>
<blockquote>
<p>DataParallel(<br>
(module): MobileNetV2(<br>
(features): Sequential(<br>
…<br>
…<br>
)<br>
)<br>
(classifier): Sequential(<br>
(0): Linear(in_features=1280, out_features=512, bias=True)<br>
(1): ReLU()<br>
(2): Dropout(p=0.2, inplace=False)<br>
(3): Linear(in_features=512, out_features=7, bias=True)<br>
)<br>
)<br>
(classifier): Identity()<br>
)</p>
</blockquote>
<p>You can see that there are two classifiers (FC layers) because I use 2 GPUs.</p>
<p>One More Important thing:</p>
<p>When I set the number of GPUs to 1, so I do not have the above problem, I get this error:</p>
<p><code>OpenBLAS Warning : Detect OpenMP Loop and this application may hang. Please rebuild the library with USE_OPENMP=1 option.</code></p>

username: ayalaa2
cooked: <p>Break your model out of DataParallel, replace the layers, and then wrap it with DataParallel again. I’m assuming you won’t need to do this enough to cause any severe overhead.</p>
<p>I’m unsure on that warning. I’m assuming this is a separate problem altogether.</p>

username: NoobCoder
cooked: <p>Thanks for you answer.</p>
<p>Actually I have though about your idea, but I wanted to do with same procedures I have done in my previous codes.</p>
<p>BTW, is there any specific code to break out from DataParallel? I was not able to find any.</p>

username: ayalaa2
cooked: <p>You can get it out like so: <code>model = model.module</code> where model was initially wrapped by DataParallel.</p>
<p>It’s kind of annoying because, depending on your code, you might have to have if statements surrounding it. (eg., if using single gpu, don’t do this)</p>

username: NoobCoder
cooked: <p>Thanks for your responses.</p>
<p>I will put the code maybe it is useful for others:</p>
<pre><code class="lang-auto"># Unwrap from DataParallel
    if torch.cuda.device_count() &gt; 1:
        model = model.module

    model.classifier = Identity()

    if torch.cuda.device_count() &gt; 1:
        print("Wrap Again with DataParallel with ", torch.cuda.device_count(),
              "GPUs!")
        model = nn.DataParallel(model)
</code></pre>