Torch.cat(list(x.values()), dim=dim), nn.DataParallel: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

Feng_Shi: <p>Hello!<br>
I am testing following code as this blog:<br>
<a href="https://github.com/dthiagarajan/technical_blog/blob/master/_notebooks/2020-03-18-Dynamic-UNet-and-PyTorch-Hooks.ipynb" rel="noopener nofollow ugc">https://github.com/dthiagarajan/technical_blog/blob/master/_notebooks/2020-03-18-Dynamic-UNet-and-PyTorch-Hooks.ipynb</a></p>
<p>It runs OK when using only single gpu. However, if using multiple gpus with  “model= nn.DataParallel(model)”, it will gives errors of:<br>
"<br>
in forward<br>
return torch.cat(list(x.values()), dim=dim)<br>
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0! (when checking arugment for argument tensors in method wrapper__cat)<br>
"<br>
I guess it may be related to “torch.cat(list(x.values()), dim=dim)”, which using python list instead of nn.ModuleList, also maybe some calls using python dict instead of  nn.ModuleDict.</p>
<p>simple change list() to nn.ModuleList won’t work.</p>
<p>anyone please help on how to solve such problem</p>

ptrblck: <p>Could you check the <code>.device</code> attribute of all tensors in <code>x</code> before this <code>torch.cat</code> operation is performed?<br>
The linked notebook is unable to render the code for some reason for me, so I’m unsure how this layer and operations are used.</p>
<p>Generally, we also recommend using <code>DistributedDataParallel</code> for better performance and since <code>nn.DataParallel</code> is in maintenance mode.</p>

Feng_Shi: <p>thks, would you posible to run following code:</p>
<p><span class="hashtag-raw">#source:</span> <a href="https://github.com/dthiagarajan/technical_blog/blob/master/_notebooks/2020-03-18-Dynamic-UNet-and-PyTorch-Hooks.ipynb" class="inline-onebox" rel="noopener nofollow ugc">technical_blog/_notebooks/2020-03-18-Dynamic-UNet-and-PyTorch-Hooks.ipynb at master · dthiagarajan/technical_blog · GitHub</a><br>
import matplotlib.pyplot as plt<br>
import numpy as np<br>
import os<br>
from PIL import Image</p>
<p>import torch<br>
import torch.nn as nn<br>
import torch.optim as optim<br>
from torch.nn import functional as F<br>
import torch.utils.model_zoo as model_zoo<br>
import torchvision.transforms.functional as tf<br>
from tqdm.notebook import tqdm</p>
<p>model_urls = {<br>
‘resnet34’: ‘<a href="https://download.pytorch.org/models/resnet34-333f7ec4.pth" rel="noopener nofollow ugc">https://download.pytorch.org/models/resnet34-333f7ec4.pth</a>’<br>
}</p>
<p><span class="hashtag-raw">#collapse_hide</span><br>
def conv3x3(in_planes, out_planes, stride=1):<br>
“”“3x3 convolution with padding”“”<br>
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,<br>
padding=1, bias=False)</p>
<p>def conv1x1(in_planes, out_planes, stride=1):<br>
“”“1x1 convolution”“”<br>
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)</p>
<p>class BasicBlock(nn.Module):<br>
expansion = 1</p>
<pre><code>def __init__(self, inplanes, planes, stride=1, downsample=None):
    super(BasicBlock, self).__init__()
    self.conv1 = conv3x3(inplanes, planes, stride)
    self.bn1 = nn.BatchNorm2d(planes)
    self.relu = nn.ReLU(inplace=True)
    self.conv2 = conv3x3(planes, planes)
    self.bn2 = nn.BatchNorm2d(planes)
    self.downsample = downsample
    self.stride = stride

def forward(self, x):
    residual = x

    out = self.conv1(x)
    out = self.bn1(out)
    out = self.relu(out)

    out = self.conv2(out)
    out = self.bn2(out)

    if self.downsample is not None:
        residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out
</code></pre>
<p>class Bottleneck(nn.Module):<br>
expansion = 4</p>
<pre><code>def __init__(self, inplanes, planes, stride=1, downsample=None):
    super(Bottleneck, self).__init__()
    self.conv1 = conv1x1(inplanes, planes)
    self.bn1 = nn.BatchNorm2d(planes)
    self.conv2 = conv3x3(planes, planes, stride)
    self.bn2 = nn.BatchNorm2d(planes)
    self.conv3 = conv1x1(planes, planes * self.expansion)
    self.bn3 = nn.BatchNorm2d(planes * self.expansion)
    self.relu = nn.ReLU(inplace=True)
    self.downsample = downsample
    self.stride = stride

def forward(self, x):
    residual = x

    out = self.conv1(x)
    out = self.bn1(out)
    out = self.relu(out)

    out = self.conv2(out)
    out = self.bn2(out)
    out = self.relu(out)

    out = self.conv3(out)
    out = self.bn3(out)

    if self.downsample is not None:
        residual = self.downsample(x)

    out += residual
    out = self.relu(out)

    return out
</code></pre>
<p>class ResNetEncoder(nn.Module):<br>
def <strong>init</strong>(self, block, layers, num_classes=1000):<br>
super(ResNetEncoder, self).<strong>init</strong>()<br>
self.inplanes = 64<br>
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,<br>
bias=False)<br>
self.bn1 = nn.BatchNorm2d(64)<br>
self.relu = nn.ReLU(inplace=True)<br>
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)<br>
self.layer0 = nn.Sequential(self.conv1, self.bn1, self.relu)<br>
self.layer1 = nn.Sequential(self.maxpool, self._make_layer(block, 64, layers[0]))<br>
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)<br>
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)<br>
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)<br>
self.out_dim = 512 * block.expansion</p>
<pre><code>    for m in self.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.constant_(m.weight, 1)
            nn.init.constant_(m.bias, 0)

def _make_layer(self, block, planes, blocks, stride=1):
    downsample = None
    if stride != 1 or self.inplanes != planes * block.expansion:
        downsample = nn.Sequential(
            conv1x1(self.inplanes, planes * block.expansion, stride),
            nn.BatchNorm2d(planes * block.expansion),
        )

    layers = []
    layers.append(block(self.inplanes, planes, stride, downsample))
    self.inplanes = planes * block.expansion
    for _ in range(1, blocks):
        layers.append(block(self.inplanes, planes))

    return nn.Sequential(*layers)

def forward(self, x):
    x = self.layer0(x)
    x = self.layer1(x)
    x = self.layer2(x)
    x = self.layer3(x)
    x = self.layer4(x)

    return x
</code></pre>
<p><span class="hashtag-raw">#collapse_show</span><br>
def resnet18(pretrained=True, **kwargs):<br>
“”“Constructs a ResNet-18 model.<br>
Args:<br>
pretrained (bool): If True, returns a model pre-trained on ImageNet<br>
“””<br>
model = ResNetEncoder(BasicBlock, [2, 2, 2, 2], **kwargs)<br>
if pretrained:<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet18’]), strict=False)<br>
return model</p>
<p>def resnet34(pretrained=True, **kwargs):<br>
“”“Constructs a ResNet-34 model.<br>
Args:<br>
pretrained (bool): If True, returns a model pre-trained on ImageNet<br>
“””<br>
model = ResNetEncoder(BasicBlock, [3, 4, 6, 3], **kwargs)<br>
if pretrained:<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet34’]), strict=False)<br>
return model</p>
<p>def resnet50(pretrained=True, **kwargs):<br>
“”“Constructs a ResNet-50 model.<br>
Args:<br>
pretrained (bool): If True, returns a model pre-trained on ImageNet<br>
“””<br>
model = ResNetEncoder(Bottleneck, [3, 4, 6, 3], **kwargs)<br>
if pretrained:<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet50’]), strict=False)<br>
return model</p>
<p>def resnet101(pretrained=True, **kwargs):<br>
“”“Constructs a ResNet-101 model.<br>
Args:<br>
pretrained (bool): If True, returns a model pre-trained on ImageNet<br>
“””<br>
model = ResNetEncoder(Bottleneck, [3, 4, 23, 3], **kwargs)<br>
if pretrained:<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet101’]), strict=False)<br>
return model</p>
<p>def resnet152(pretrained=True, **kwargs):<br>
“”“Constructs a ResNet-152 model.<br>
Args:<br>
pretrained (bool): If True, returns a model pre-trained on ImageNet<br>
“””<br>
model = ResNetEncoder(Bottleneck, [3, 8, 36, 3], **kwargs)<br>
if pretrained:<br>
model.load_state_dict(model_zoo.load_url(model_urls[‘resnet152’]), strict=False)<br>
return model</p>
<p>class ConvLayer(nn.Module):<br>
def <strong>init</strong>(self, num_inputs, num_filters, bn=True, kernel_size=3, stride=1,<br>
padding=None, transpose=False, dilation=1):<br>
super(ConvLayer, self).<strong>init</strong>()<br>
if padding is None:<br>
padding = (kernel_size-1)//2 if transpose is not None else 0<br>
if transpose:<br>
self.layer = nn.ConvTranspose2d(num_inputs, num_filters, kernel_size=kernel_size,<br>
stride=stride, padding=padding, dilation=dilation)<br>
else:<br>
self.layer = nn.Conv2d(num_inputs, num_filters, kernel_size=kernel_size,<br>
stride=stride, padding=padding)<br>
nn.init.kaiming_uniform_(self.layer.weight, a=np.sqrt(5))<br>
self.bn_layer = nn.BatchNorm2d(num_filters) if bn else None</p>
<pre><code>def forward(self, x):
    out = self.layer(x)
    out = F.relu(out)
    return out if self.bn_layer is None else self.bn_layer(out)
</code></pre>
<p>class ConcatLayer(nn.Module):<br>
def forward(self, x, dim=1):<br>
return torch.cat(list(x.values()), dim=dim)</p>
<p>class LambdaLayer(nn.Module):<br>
def <strong>init</strong>(self, f):<br>
super(LambdaLayer, self).<strong>init</strong>()<br>
self.f = f</p>
<pre><code>def forward(self, x):
    return self.f(x)
</code></pre>
<p>def upconv2x2(inplanes, outplanes, size=None, stride=1):<br>
if size is not None:<br>
return [<br>
ConvLayer(inplanes, outplanes, kernel_size=2, dilation=2, stride=stride),<br>
nn.Upsample(size=size, mode=‘bilinear’, align_corners=True)<br>
]<br>
else:<br>
return [<br>
ConvLayer(inplanes, outplanes, kernel_size=2, dilation=2, stride=stride),<br>
nn.Upsample(scale_factor=2, mode=‘bilinear’, align_corners=True)<br>
]</p>
<p><span class="hashtag-raw">#collapse_show</span><br>
class DecoderConnect(nn.Module):<br>
def <strong>init</strong>(self, inplanes, output_size):<br>
super(DecoderConnect, self).<strong>init</strong>()<br>
self.bottom_process = nn.Sequential(<br>
ConvLayer(inplanes, inplanes * 2, kernel_size=3),<br>
ConvLayer(inplanes * 2, inplanes * 2, kernel_size=3),<br>
*upconv2x2(inplanes * 2, inplanes, size=output_size)<br>
)<br>
self.concat_process = nn.Sequential(<br>
ConcatLayer(),<br>
ConvLayer(inplanes * 2, inplanes * 2, kernel_size=1),<br>
ConvLayer(inplanes * 2, inplanes, kernel_size=3),<br>
ConvLayer(inplanes, inplanes, kernel_size=3)<br>
)</p>
<pre><code>def forward(self, x):
    decoder_input = self.bottom_process(x)
    return self.concat_process({0: x, 1: decoder_input})
</code></pre>
<p><span class="hashtag-raw">#collapse_show</span><br>
class DynamicUNet(nn.Module):<br>
def <strong>init</strong>(self, encoder, input_size=(224, 224), num_output_channels=None, verbose=0):<br>
super(DynamicUNet, self).<strong>init</strong>()<br>
self.encoder = encoder<br>
self.verbose = verbose<br>
self.input_size = input_size<br>
self.num_input_channels = 3  # This must be 3 because we’re using a ResNet encoder<br>
self.num_output_channels = num_output_channels</p>
<pre><code>    self.decoder = self.setup_decoder()
    
def forward(self, x):
    encoder_outputs = []
    def encoder_output_hook(self, input, output):
        encoder_outputs.append(output)

    handles = [
        child.register_forward_hook(encoder_output_hook) for name, child in self.encoder.named_children()
        if name.startswith('layer')
    ]

    try:
        self.encoder(x)
    finally:
        if self.verbose &gt;= 1:
            print("Removing all forward handles")
        for handle in handles:
            handle.remove()

    prev_output = None
    for reo, rdl in zip(reversed(encoder_outputs), self.decoder):
        if prev_output is not None:
            prev_output = rdl({0: reo, 1: prev_output})
        else:
            prev_output = rdl(reo)
    return prev_output
            
def setup_decoder(self):
    input_sizes = []
    output_sizes = []
    def shape_hook(self, input, output):
        input_sizes.append(input[0].shape)
        output_sizes.append(output.shape)

    handles = [
        child.register_forward_hook(shape_hook) for name, child in self.encoder.named_children()
        if name.startswith('layer')
    ]    

    self.encoder.eval()
    test_input = torch.randn(1, self.num_input_channels, *self.input_size)
    try:
        self.encoder(test_input)
    finally:
        if self.verbose &gt;= 1:
            print("Removing all shape hook handles")
        for handle in handles:
            handle.remove()
    decoder = self.construct_decoder(input_sizes, output_sizes, num_output_channels=self.num_output_channels)
    return decoder
    
def construct_decoder(self, input_sizes, output_sizes, num_output_channels=None):
    decoder_layers = []
    for layer_index, (input_size, output_size) in enumerate(zip(input_sizes, output_sizes)):
        upsampling_size_factor = int(input_size[-1] / output_size[-1])
        upsampling_channel_factor = input_size[-3] / output_size[-3]
        next_layer = []
        bs, c, h, w = input_size
        ops = []
        if layer_index == len(input_sizes) - 1:
            last_layer_ops = DecoderConnect(output_size[-3], output_size[2:])
            last_layer_ops_input = torch.randn(*output_size)
            last_layer_concat_ops_output = last_layer_ops(last_layer_ops_input)
            next_layer.extend([last_layer_ops])
            if upsampling_size_factor &gt; 1 or upsampling_channel_factor != 1:
                last_layer_concat_upconv_op = upconv2x2(output_size[-3], input_size[-3], size=input_size[2:])
                last_layer_concat_upconv_op_output = nn.Sequential(*last_layer_concat_upconv_op)(
                    last_layer_concat_ops_output
                )
                next_layer.extend(last_layer_concat_upconv_op)
        elif layer_index == 0:
            first_layer_concat_ops = [
                ConcatLayer(),
                ConvLayer(output_size[-3] * 2, output_size[-3] * 2, kernel_size=1),
                *upconv2x2(
                    output_size[-3] * 2,
                    output_size[-3],
                    size=[dim * upsampling_size_factor for dim in output_size[2:]]
                ),
                ConvLayer(output_size[-3], output_size[-3], kernel_size=3),
                ConvLayer(
                    output_size[-3],
                    input_size[-3] if self.num_output_channels is None else self.num_output_channels,
                    kernel_size=1
                ),
            ]
            first_layer_concat_ops_output = nn.Sequential(*first_layer_concat_ops)(
                {0: torch.randn(*output_size), 1: torch.randn(*output_size)}
            )
            next_layer.extend(first_layer_concat_ops)
        else:
            middle_layer_concat_ops = [
                ConcatLayer(),
                ConvLayer(output_size[-3] * 2, output_size[-3] * 2, kernel_size=1),
                ConvLayer(output_size[-3] * 2, output_size[-3], kernel_size=3),
                ConvLayer(output_size[-3], output_size[-3], kernel_size=3)
            ]
            middle_layer_concat_ops_output = nn.Sequential(*middle_layer_concat_ops)(
                {0: torch.randn(*output_size), 1: torch.randn(*output_size)}
            )
            next_layer.extend(middle_layer_concat_ops)
            if upsampling_size_factor &gt; 1 or upsampling_channel_factor != 1:
                middle_layer_concat_upconv_op = upconv2x2(output_size[-3], input_size[-3], size=input_size[2:])
                middle_layer_concat_upconv_op_output = nn.Sequential(*middle_layer_concat_upconv_op)(
                    middle_layer_concat_ops_output
                )
                next_layer.extend(middle_layer_concat_upconv_op)
        decoder_layers.append(nn.Sequential(*next_layer))
    return nn.ModuleList(reversed(decoder_layers))
</code></pre>
<p>model = DynamicUNet(resnet34(), num_output_channels=3, input_size=(256,256))<br>
<span class="hashtag-raw">#if</span> torch.cuda.is_available():</p>
<h1><a name="model-modelcuda-1" class="anchor" href="#model-modelcuda-1"></a>model = model.cuda()</h1>
<p><span class="hashtag-raw">#decoder_parameters</span> = [item for module in model.decoder for item in module.parameters()]<br>
<span class="hashtag-raw">#optimizer</span> = optim.AdamW(decoder_parameters)  # Only training the decoder for now</p>
<p>model.train()<br>
if torch.cuda.device_count() &gt; 1:<br>
print(“Let’s use”, torch.cuda.device_count(), “GPUs!”)<br>
# dim = 0 [30, xxx] → [10, …], [10, …], [10, …] on 3 GPUs<br>
model= nn.DataParallel(model)<br>
model = model.cuda()<br>
inputs = torch.randn(8, 3, 256, 256).cuda()</p>
<p>out = model(inputs)</p>