Network partitioning

Mohammad_Javadi: <p>I tried to train real-esrgan on my dataset with a single GPU, but couldn’t due to GPU memory limitation.<br>
However, the machine I am running this model on has multiple GPUs, so I considered partitioning the model and training sub-networks in multiple GPUs. Here is the code I used:</p>
<pre><code class="lang-auto">class SRResNet(nn.Module):

    def __init__(self, img_feat = 1, n_feats = 64, kernel_size = 3, num_block = 16, act = nn.PReLU(), scale=4):
        super(SRResNet, self).__init__()

        self.conv01 = conv(in_channel = img_feat, out_channel = n_feats, kernel_size = 9, BN = False, act = act).to('cuda:0')

        resblocks = [ResBlock(channels = n_feats, kernel_size = 3, act = act) for _ in range(int(num_block/2))]
        self.body1 = nn.Sequential(*resblocks).to('cuda:0')
        self.body2 = nn.Sequential(*resblocks).to('cuda:1')

        self.conv02 = conv(in_channel = n_feats, out_channel = n_feats, kernel_size = 3, BN = True, act = None).to('cuda:1')

        self.last_conv = conv(in_channel = n_feats, out_channel = img_feat, kernel_size = 3, BN = False, act = nn.Tanh()).to('cuda:1')

    def forward(self, x):

        x = self.conv01(x)
        _skip_connection = x.to('cuda:0')

        x = self.body1(x)
        x = self.body2(x.to('cuda:1'))
        x = self.conv02(x)
        copy_connection = _skip_connection.to('cuda:1')
        feat = x + copy_connection

        x = self.last_conv(feat)

        return x
</code></pre>
<p>I load SRResNet like:</p>
<pre><code class="lang-auto">model = SRResNet()
</code></pre>
<p>I don’t move the network to any specific device since I specified each layer’s device in the network implementation.</p>
<p>I also load my input tensor on gpu0 and my ground truth on gpu1.</p>
<p>I get the following error in the forward function:</p>
<pre><code class="lang-auto">RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument weight in method wrapper__prelu)
</code></pre>
<p>Any ideas on how to train this network on multiple GPUs?</p>

ptrblck: <p>The error points to <code>PReLU</code> which is the default <code>act</code> argument and used in:</p>
<pre><code class="lang-python">self.conv01 = conv(in_channel = img_feat, out_channel = n_feats, kernel_size = 9, BN = False, act = act).to('cuda:0')
</code></pre>
<p>as well as:</p>
<pre><code class="lang-python">resblocks = [ResBlock(channels = n_feats, kernel_size = 3, act = act) for _ in range(int(num_block/2))]
        self.body1 = nn.Sequential(*resblocks).to('cuda:0')
        self.body2 = nn.Sequential(*resblocks).to('cuda:1')
</code></pre>
<p>Based on this code the second one looks wrong as you are using references to <code>resblock</code> while I assume you want to create different instances:</p>
<pre><code class="lang-python">resblocks = [nn.Linear(10, 10)]
body1 = nn.Sequential(*resblocks).to('cpu')
print(body1[0].weight.device)
# cpu
print(resblocks[0].weight.device)
# cpu

body2 = nn.Sequential(*resblocks).to('cuda')
print(body1[0].weight.device)
# cuda:0
print(body2[0].weight.device)
# cuda:0
print(resblocks[0].weight.device)
# cuda:0
</code></pre>
<p>Here you can see that both <code>self.bodyX</code> modules will reuse the same <code>resblock</code> and the second <code>to()</code> operation will thus also move <code>self.body1</code> to <code>cuda:1</code>.</p>