DataParallel donnot support customed layer

shaoming20798: <p>Hi, I have got a problem abount multi-gpus training. The code is from <code>https://github.com/chaoqichen/HTCN</code>. Here is the key code:</p>
<pre data-code-wrap="python"><code class="lang-python">class RandomLayer(nn.Module):

    def __init__(self, input_dim_list=[], output_dim=1024):
        super(RandomLayer, self).__init__()
        self.input_num = len(input_dim_list)# 2
        self.output_dim = output_dim
        self.random_matrix = [torch.rand(input_dim_list[i], output_dim) for i in range(self.input_num)]

    def forward(self, input_list):
        return_list = [torch.mm(input_list[i], self.random_matrix[i]) for i in range(self.input_num)]
        return_tensor = return_list[0] / math.pow(float(self.output_dim), 1.0 / len(return_list))
        for single in return_list[1:]:
            return_tensor = torch.mul(return_tensor, single)
        return return_tensor

    def cuda(self):
        super(RandomLayer, self).cuda()
        self.random_matrix = [val.cuda() for val in self.random_matrix]

class vgg16(_fasterRCNN):

  def __init__(self, classes, pretrained=False, class_agnostic=False,lc=False,gc=False, la_attention = False, mid_attention = False):
    self.model_path = cfg.VGG_PATH
    self.dout_base_model = 512
    self.pretrained = pretrained
    self.class_agnostic = class_agnostic
    self.lc = lc
    self.gc = gc

    _fasterRCNN.__init__(self, classes, class_agnostic,lc,gc, la_attention, mid_attention)

  def _init_modules(self):
    vgg = models.vgg16()
    if self.pretrained:
        print("Loading pretrained weights from %s" %(self.model_path))
        state_dict = torch.load(self.model_path)
        vgg.load_state_dict({k:v for k,v in state_dict.items() if k in vgg.state_dict()})

    vgg.classifier = nn.Sequential(*list(vgg.classifier._modules.values())[:-1])

    # not using the last maxpool layer
    #print(vgg.features)
    self.RCNN_base1 = nn.Sequential(*list(vgg.features._modules.values())[:14])
    self.RCNN_base2 = nn.Sequential(*list(vgg.features._modules.values())[14:21])
    self.RCNN_base3 = nn.Sequential(*list(vgg.features._modules.values())[21:-1])
    #print(self.RCNN_base1)
    #print(self.RCNN_base2)
    self.netD = netD(context=self.gc)
    self.netD_pixel = netD_pixel(context=self.lc)
    self.netD_mid = netD_mid(context=self.gc)
    feat_d = 4096
    feat_d2 = 384
    feat_d3 = 2048

    self.RandomLayer = RandomLayer([feat_d, feat_d2], feat_d3)
    self.RandomLayer.cuda()
    # Fix the layers before conv3:
    self.netD_da = netD_da(feat_d3)

    for layer in range(10):
      for p in self.RCNN_base1[layer].parameters(): p.requires_grad = False

    # self.RCNN_base = _RCNN_base(vgg.features, self.classes, self.dout_base_model)

    self.RCNN_top = vgg.classifier

    self.RCNN_cls_score = nn.Linear(feat_d+feat_d2, self.n_classes)
    if self.class_agnostic:
        self.RCNN_bbox_pred = nn.Linear(feat_d+feat_d2, 4)
    else:
        self.RCNN_bbox_pred = nn.Linear(feat_d+feat_d2, 4 * self.n_classes)


  def _head_to_tail(self, pool5):
    
    pool5_flat = pool5.view(pool5.size(0), -1)
    fc7 = self.RCNN_top(pool5_flat)

    return fc7
</code></pre>
<p><strong>RandomLayer</strong> cannot be assigned to multi devices using the <code>.cuda()</code> method using the code <code>self.RandomLayer.cuda()</code>. The error message is as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">return_list = [torch.mm(input_list[i], self.random_matrix[i]) for i in range(self.input_num)]
</code></pre>
<blockquote>
<p>RuntimeError: Expected tensor for ‘out’ to have the same device as tensor for argument <span class="hashtag-raw">#3</span> ‘mat2’; but device 1 dose not equal 0(while checking arguments for addmm)</p>
</blockquote>
<p>And I print the device of <code>input_list[i]</code> and <code>self.random_matrix[i]</code> when using two gpus, the result is as follows (<code>self.input_num == 2</code>):</p>
<blockquote>
<p>gpu id 0:<br>
input_list0: cuda:0 self.random_matrix0: cuda:0<br>
input_list1: cuda:0 self.random_matrix1: cuda:0<br>
gpu id 1:<br>
input_list0: cuda:1 self.random_matrix0: cuda:0<br>
input_list1: cuda:1 self.random_matrix1: cuda:0</p>
</blockquote>
<p>Is there any advice to fix the problem?</p>

shaoming20798: <p>I changed the ‘RandomLayer’ code and it’s working with 2 gpus, but still failed with 3 or 4 gpus. I donnot know why, so could you please have a loot at my question? In <code>https://github.com/pytorch/pytorch/issues/50087</code>,<a class="mention" href="/u/ngimel">@ngimel</a> told me that you cannot override <code>.cuda()</code> method this way, it’s not supported, you need to override <code>_apply</code> method instead. But I don’t know how to implemented it.<br>
<a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Here is my modified <code>RandomLyaer</code> codes, because len(self.random_matrix) is actually 2 in my project, so I simply split ‘self.random_matrix[2]’ to 2 scalars. 2 gpus is working, 3 or more gpus will be failed.</p>
<pre><code class="lang-python">class RandomLayer(nn.Module):
    def __init__(self, batch_size, input_dim_list=[], output_dim=1024):
        super(RandomLayer, self).__init__()
        self.input_num = len(input_dim_list)# 2
        self.output_dim = output_dim
        self.random_matrix0 = torch.nn.Parameter(torch.rand(input_dim_list[0], output_dim))
        self.random_matrix1 = torch.nn.Parameter(torch.rand(input_dim_list[1], output_dim))

    def forward(self, input_list):
        assert(len(input_list) == 2)
        return_list0 = torch.mm(input_list[0], self.random_matrix0)
        return_list1 = torch.mm(input_list[1], self.random_matrix1)
        return_tensor = return_list0 / math.pow(float(self.output_dim), 1.0 / 2)
        return_tensor = torch.mul(return_tensor, return_list1)
        return return_tensor
</code></pre>

ptrblck: <p>I’m not sure why you want to override the <code>cuda</code> method, but guess you are running into another device mismatch with <code>self.random_matrix</code>, since it’s initialized as a plain <code>list</code>, and tried to fix it with the custom <code>cuda</code> method.</p>
<p>To write device-agnostic code remove all <code>cuda()</code> calls and make sure all modules, parameters, and buffers are properly registered.<br>
E.g. <code>self.random_matrix</code> won’t register the tensors and you would either have to use e.g. <code>nn.ParameterList</code> or register all tensors via <code>self.register_buffer</code>.</p>
<p>This will make sure that calling <code>random_layer.cuda()</code> pushes all internal submodules etc. to the defined device (which is used in <code>nn.DataParallel</code>). The updated module looks generally alright and works fine in this code snippet on 8 GPUs:</p>
<pre><code class="lang-python">class RandomLayer(nn.Module):
    def __init__(self, batch_size, input_dim_list=[], output_dim=1024):
        super(RandomLayer, self).__init__()
        self.input_num = len(input_dim_list)# 2
        self.output_dim = output_dim
        self.random_matrix0 = torch.nn.Parameter(torch.rand(input_dim_list[0], output_dim))
        self.random_matrix1 = torch.nn.Parameter(torch.rand(input_dim_list[1], output_dim))

    def forward(self, x1, x2):
        return_list0 = torch.mm(x1, self.random_matrix0)
        return_list1 = torch.mm(x2, self.random_matrix1)
        return_tensor = return_list0 / math.pow(float(self.output_dim), 1.0 / 2)
        return_tensor = torch.mul(return_tensor, return_list1)
        return return_tensor


class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.rand = RandomLayer(1, [1, 1])

    def forward(self, x1, x2):
        print(x1.device)
        x = self.rand(x1, x2)
        return x

model = MyModule()
model = nn.DataParallel(model).cuda()
x1 = torch.randn(8, 1).cuda()
x2 = torch.randn(8, 1).cuda()
out = model(x1, x2)
print(out.device)
</code></pre>
<p>Note that <code>nn.DataParallel</code> will split the input batch in <code>dim0</code> and send each chunk to the corresponding GPU. In your current code snippet you are trying to index <code>input_list</code> in <code>dim0</code>, which also seems to be wrong, as it would index it in the (split) batch dimension.<br>
You could thus either pass multiple inputs (as done in my code snippet) or stack and index the input in another dimension.</p>

shaoming20798: <p>Hi ,ptrblck. Thank you for your kindly reply. I try to use nn.ParameterList in multi-gpus mode in your coe snippet, but it still doesn’t work. Here is the source code named ‘dp.py’, single gpu is OK.</p>
<pre><code class="lang-python">#!/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import math

feat_d = 16
feat_d2 = 8
feat_d3 = 8

class RandomLayer(nn.Module):
    def __init__(self, input_dim_list=[], output_dim=1024):
        super(RandomLayer, self).__init__()
        self.input_num = len(input_dim_list)# 2
        self.output_dim = output_dim
        self.random_matrix = torch.nn.ParameterList([torch.nn.Parameter(torch.rand(input_dim_list[i], output_dim)) for i in range(self.input_num)])
        print('*****')
        print(len(self.random_matrix))
        print(self.random_matrix)

    def forward(self, input_list):
        print('&amp;&amp;&amp;&amp;&amp;')
        print(len(self.random_matrix))
        print(self.random_matrix)
        return_list = [torch.mm(input_list[i], self.random_matrix[i]) for i in range(self.input_num)]
        return_tensor = return_list[0] / math.pow(float(self.output_dim), 1.0 / 2)
        for single in return_list[1:]:
            return_tensor = torch.mul(return_tensor, single)
        return return_tensor


class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.rand = RandomLayer([feat_d, feat_d2], feat_d3)

    def forward(self, x1, x2):
        print(x1.device)
        x = self.rand([x1, x2])
        return x

model = MyModule()
model = nn.DataParallel(model).cuda()
x1 = torch.randn(4, feat_d).cuda()
x2 = torch.randn(4, feat_d2).cuda()
out = model(x1, x2)
print(out.device)
</code></pre>
<p>when I running that code with the command ‘CUDA_VISIBLE_DEVICES=4,5 python dp.py’, ‘len(self.random_matrix)’ is 0. I don’t know why, so could you have a look at this?</p>

ptrblck: <p><code>len(self.random_matrix)</code> is <code>0</code>, since you are not passing <code>input_dim_list</code> and use the default value, which is an empty <code>list</code>.<br>
This also means that <code>self.input_num = len(input_dim_list)</code> will be <code>0</code> and thus:</p>
<pre><code class="lang-python">self.random_matrix = torch.nn.ParameterList([torch.nn.Parameter(torch.rand(input_dim_list[i], output_dim)) for i in range(self.input_num)])
</code></pre>
<p>will be empty.</p>

shaoming20798: <p>Hi, ptrblck.<br>
I think I have passed a <code>list</code> to <code>input_dim_list</code> in that code which is as follows:<br>
<code>self.rand = RandomLayer([feat_d, feat_d2], feat_d3)</code>. And that code could work in single gpu, giving <code>len(self.random_matrix)</code> equals to 2 correctly. But 2 gpu will give a wrong result, with <code>len(self.random_matrix)</code> equals to 0.<br>
And I found that in 2 gpus mode, in the class <code>RandomLayer</code> <code>__init__</code> method, <code>len(self.random_matrix)</code> equals to 2. However, in  <code>RandomLayer</code> <code>forward</code> method, <code>len(self.random_matrix</code> equals to 0.</p>

ptrblck: <p>You are right and I missed the constructor inside <code>MyModule</code>.<br>
You should get a warning using <code>nn.ParameterList</code> and <code>nn.DataParallel</code>, which is not supported in this setup, so you would need to fallback to registering the parameters separately:</p>
<pre><code class="lang-python">UserWarning: nn.ParameterList is being used with DataParallel but this is not supported. This list will appear empty for the models replicated on each GPU except the original one.
</code></pre>