PyTorch Model throwing error for list of layers

Admin 1
Viewed 270

I have designed the following torch model with 2 conv2d layers. It works without any error.

import torch.nn as nn
from torchsummary import summary

class mini_unet(nn.Module):
    def __init__(self):
        super(mini_unet, self).__init__()
        self.c1 = nn.Conv2d(1, 1, 3, padding = 1)
        self.r1 = nn.ReLU()
        self.c2 = nn.Conv2d(1, 1, 3, padding = 1)
        self.r2 = nn.ReLU()

    def forward(self, x):
        x = self.c1(x)
        x = self.r1(x)
        x = self.c2(x)
        x = self.r2(x)
        return x

a = mini_unet().cuda()

print(a)

But, let's say I have too many layers, I don't want to explicitly write each of them in the forward function. So, I used a list to automate it like below.

import torch.nn as nn
from torchsummary import summary

class mini_unet2(nn.Module):
    def __init__(self):
        super(mini_unet2, self).__init__()
        self.layers = [nn.Conv2d(1, 1, 3, padding = 1),
        nn.ReLU(),
        nn.Conv2d(1, 1, 3, padding = 1),
        nn.ReLU()]

    def forward(self, x):
        for l in self.layers:
            x = l(x)
        return x
    
a2 = mini_unet2().cuda()
print(a2)
summary(a2, (1,4,4))

This gives me the following error which is strange, I have used cuda() why it doesn't work?

RuntimeError                              Traceback (most recent call last)
<ipython-input-36-1d71e75b96e0> in <module>
     17 a2 = mini_unet2().cuda()
     18 print(a2)
---> 19 summary(a2, (1,4,4))

~/anaconda3/envs/torch/lib/python3.6/site-packages/torchsummary/torchsummary.py in summary(model, input_size, batch_size, device)
     70     # make a forward pass
     71     # print(x.shape)
---> 72     model(*x)
     73 
     74     # remove these hooks

~/anaconda3/envs/torch/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

<ipython-input-36-1d71e75b96e0> in forward(self, x)
     12     def forward(self, x):
     13         for l in self.layers:
---> 14             x = l(x)
     15         return x
     16 

~/anaconda3/envs/torch/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    487             result = self._slow_forward(*input, **kwargs)
    488         else:
--> 489             result = self.forward(*input, **kwargs)
    490         for hook in self._forward_hooks.values():
    491             hook_result = hook(self, input, result)

~/anaconda3/envs/torch/lib/python3.6/site-packages/torch/nn/modules/conv.py in forward(self, input)
    318     def forward(self, input):
    319         return F.conv2d(input, self.weight, self.bias, self.stride,
--> 320                         self.padding, self.dilation, self.groups)
    321 
    322 

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
torch pytorch python-3.x
1 Answers

The error is maybe a little counter-intuitive but the error originates from you using python lists for the layers.

From the documentation, you need to use torch.nn.ModuleList to contain the submodules, not a python list.

So, just changing the list with nn.Modulelist(list) will solve the error.

import torch.nn as nn
from torchsummary import summary

class mini_unet2(nn.Module):
    def __init__(self):
        super(mini_unet2, self).__init__()
        self.layers = nn.ModuleList([nn.Conv2d(1, 1, 3, padding = 1),
        nn.ReLU(),
        nn.Conv2d(1, 1, 3, padding = 1),
        nn.ReLU()])

    def forward(self, x):
        for l in self.layers:
            x = l(x)
        return x

a2 = mini_unet2().cuda()
print(a2)
summary(a2, (1,4,4))
user_46220460
Related