How can I use torch.nn.DataParallel while I'm doing transfer learning?

username: pouya.ahmadvand
cooked: <p>I’m trying to do my training on multiple GPUs by using the following code (the latest pytorch version):</p>
<pre><code class="lang-auto">from torchvision import models
model = model.vgg16(pretrained=True)
model.classifier._modules['6'] = torch.nn.Linear(4096, 10)
self.model = torch.nn.DataParallel(model, device_ids=[0,1,2]).cuda()
self.model = model.to(f'cuda:0')
...
def forward(self, input_data):
    output = self.model.forward(input_data)
</code></pre>
<p>I get this error when I call <strong>self.model.forward(input_data)</strong> :</p>
<pre><code class="lang-auto">  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torchvision/models/vgg.py", line 43, in forward
    x = self.features(x)
  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torch/nn/modules/container.py", line 100, in forward
    input = module(input)
  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 345, in forward
    return self.conv2d_forward(input, self.weight)
  File "/home/poahmadvand/py3env/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 342, in conv2d_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 0 does not equal 1 (while checking arguments for cudnn_convolution)
</code></pre>
<p>How can I fix this error? thanks.</p>

username: pritamdamania87
cooked: <p>I can’t seem to reproduce the issue that you’re seeing. This code works fine on PyTorch 1.6:</p>
<pre><code class="lang-auto">from torchvision import models
import torch
model = models.vgg16(pretrained=True)
model.classifier._modules['6'] = torch.nn.Linear(4096, 10)
model = torch.nn.DataParallel(model, device_ids=[0,1,2]).cuda()
model = model.to(f'cuda:0')
input_data = torch.rand(10, 3, 225, 225)
model(input_data)
</code></pre>
<p>Am I missing something here?</p>

username: pouya.ahmadvand
cooked: <p>Thanks, Do you have multiple GPUs on your machine?</p>

username: pritamdamania87
cooked: <p>Yes, I’m trying this on an 8 GPU machine.</p>

username: pritamdamania87
cooked: <p><a class="mention" href="/u/pouya.ahmadvand">@pouya.ahmadvand</a> Do you run into the same error even if you run the code I pasted above?</p>

username: pouya.ahmadvand
cooked: <p>Thanks, I’ll try it tomorrow and let you know.</p>

username: pouya.ahmadvand
cooked: <p><a class="mention" href="/u/pritamdamania87">@pritamdamania87</a> Thanks, it now works. There problem was that I used a function to fetch the number of GPUs with the highest free memory available, and this function returns an array like [5, 6, 7]. Then:</p>
<pre><code class="lang-auto">selected_gpus = [5, 6, 7]
model = torch.nn.DataParallel(model, device_ids=selected_gpus).cuda()   
model = model.to(f'cuda:{selected_gpus[0]}')
</code></pre>
<p>which gives me that error. now I’m using the return array of gpu selector and set the system variable CUDA_VISIBLE_DEVICES by it and then:</p>
<pre><code class="lang-auto">selected_gpus = [5, 6, 7]
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(str(x) for x in selected_gpus)
model = torch.nn.DataParallel(model, device_ids=range(0,len(selected_gpus)).cuda()
model = model.to(f'cuda:0')
</code></pre>
<p>It works fine now.<br>
Thanks!</p>