Device mismatch error with nn.DataParallel

username: sshkhr
cooked: <p>I’m training a modified ResNet on multiple GPUs. Here is the residual block for the ResNet:</p>
<pre><code class="lang-auto">class BasicBlockWOutput(nn.Module):
    expansion = 1

    def __init__(self, in_channels, channels, params, stride=1):
        super(BasicBlockWOutput, self).__init__()
        add_output = params[0]
        num_classes = params[1]
        input_size = params[2]
        self.output_id = params[3]

        self.depth = 2

        layers = nn.ModuleList()

        conv_layer = []
        conv_layer.append(nn.Conv2d(in_channels, channels, kernel_size=3, stride=stride, padding=1, bias=False))
        conv_layer.append(nn.BatchNorm2d(channels))
        conv_layer.append(nn.ReLU())
        conv_layer.append(nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=False))
        conv_layer.append(nn.BatchNorm2d(channels))

        layers.append(nn.Sequential(*conv_layer))

        shortcut = nn.Sequential()

        if stride != 1 or in_channels != self.expansion*channels:
            shortcut = nn.Sequential(
                nn.Conv2d(in_channels, self.expansion*channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*channels)
            )

        layers.append(shortcut)
        layers.append(nn.ReLU())

        self.layers = layers

        if add_output:
            self.output = af.InternalClassifier(input_size, self.expansion*channels, num_classes) 
            self.no_output = False

        else:
            self.output = None
            self.forward = self.only_forward
            self.no_output = True
            
    def forward(self, x):
        fwd = self.layers[0](x) # conv layers
        fwd = fwd + self.layers[1](x) # shortcut
        return self.layers[2](fwd), 1, self.output(fwd)         # output layers for this module
    
    def only_output(self, x):
        fwd = self.layers[0](x) # conv layers
        fwd = fwd + self.layers[1](x) # shortcut
        fwd = self.layers[2](fwd) # activation
        out = self.output(fwd)         # output layers for this module
        return out
    
    def only_forward(self, x):
        fwd = self.layers[0](x) # conv layers
        fwd = fwd + self.layers[1](x) # shortcut
        return self.layers[2](fwd), 0, None # activation
</code></pre>
<p>When I’m running it on multiple GPUs I’m having device mismatch issues:</p>
<blockquote>
<p>RuntimeError: Caught RuntimeError in replica 1 on device 1.<br>
Original Traceback (most recent call last):<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py”, line 60, in _worker<br>
output = module(*input, **kwargs)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 722, in _call_impl<br>
result = self.forward(*input, **kwargs)<br>
File “/export/mlrg/sshekhar/XAI/Shallow-Deep-Networks-gpub/architectures/SDNs/ResNet_SDN.py”, line 159, in forward<br>
fwd, is_output, output = layer(fwd)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 722, in _call_impl<br>
result = self.forward(*input, **kwargs)<br>
File “/export/mlrg/sshekhar/XAI/Shallow-Deep-Networks-gpub/architectures/SDNs/ResNet_SDN.py”, line 68, in only_forward<br>
fwd = self.layers<a>0</a> # conv layers<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 722, in _call_impl<br>
result = self.forward(*input, **kwargs)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/container.py”, line 117, in forward<br>
input = module(input)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py”, line 722, in _call_impl<br>
result = self.forward(*input, **kwargs)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/conv.py”, line 419, in forward<br>
return self._conv_forward(input, self.weight)<br>
File “/export/mlrg/sshekhar/anaconda3/lib/python3.7/site-packages/torch/nn/modules/conv.py”, line 416, in _conv_forward<br>
self.padding, self.dilation, self.groups)<br>
RuntimeError: Expected tensor for argument <span class="hashtag-raw">#1</span> ‘input’ to have the same device as tensor for argument <span class="hashtag-raw">#2</span> ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)</p>
</blockquote>

username: osalpekar
cooked: <p>It seems like your input is on GPU 1, but your network is on GPU 0. From the error trace, it seems like the issue stems from how the <code>only_forward</code> function is defined and passed as a reference to <code>self.forward</code>.  This looks similar to the issue described in <a href="https://github.com/pytorch/pytorch/issues/8637" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/8637</a> - check out the suggestion by Ssnl about why duplicating non-tensor objects in DataParallel could lead to these device mismatch errors.</p>

username: sshkhr
cooked: <p>So I was able to resolve my issue based on Omkar’s suggestion to look at <a href="https://github.com/pytorch/pytorch/issues/8637" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/8637</a></p>
<p>I changed the following: Instead of binding <code>forward</code> to the <code>only_forward</code> method, I directly call it inside the <code>forward</code> method:</p>
<pre><code class="lang-auto">        if add_output:
            self.output = af.InternalClassifier(input_size, self.expansion*channels, num_classes) 
            self.no_output = False
        else:
            self.output = None
            #self.forward = self.only_forward
            self.no_output = True
            
    def forward(self, x):

        if self.no_output:
            return self.only_forward(x)        
        else:
            fwd = self.layers[0](x) # conv layers
            fwd = fwd + self.layers[1](x) # shortcut
            return self.layers[2](fwd), 1, self.output(fwd)         # output layers for this module

    
    def only_output(self, x):
        fwd = self.layers[0](x) # conv layers
        fwd = fwd + self.layers[1](x) # shortcut
        fwd = self.layers[2](fwd) # activation
        out = self.output(fwd)         # output layers for this module
        return out
    
    def only_forward(self, x):
        fwd = self.layers[0](x) # conv layers
        fwd = fwd + self.layers[1](x) # shortcut
        return self.layers[2](fwd), 0, None # activation
</code></pre>