Creating new layer

vision
1

Hi,
I am trying to re-implement a channel attention module in order to get my hands-on dealing with pytorch and how to create my custom layers.

class CAM_infer(nn.Module):
    """Channel Attention module"""
    
    def __init__(self, ratio= 16):
        super(CAM_infer, self).__init__()
        # self.channels = channels
        self.ratio =  ratio

        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.maxpool = nn.AdaptiveMaxPool2d((1,1))

        # self.h1 = nn.LazyLinear(channels//self.ratio, bias=True)
        # self.h2 = nn.LazyLinear(channels, bias= True)

        # nn.init.xavier_uniform_(self.h1.weight)
        # nn.init.xavier_uniform_(self.h2.weight)

        self.act = nn.Sigmoid()

        

    def forward(self, x):

        no_channels = x.size(1)
        h1 = nn.LazyLinear(no_channels//self.ratio, bias=True)
        h2 = nn.LazyLinear(no_channels, bias= True)

        x1 = self.avgpool(x)
        x2 = self.maxpool(x)

        x1 = x1.view(-1, 1 ,no_channels)
        x2 = x2.view(-1, 1 ,no_channels)

        x1 = nn.functional.relu(h1(x1))
        x1 = h2(x1)

        x2 = nn.functional.relu(h1(x2))
        x2 = h2(x2)

        x_new = torch.add(x1, x2)
        mc = self.act(x_new)
        mc = mc.view(-1, no_channels * 1)

        # x = torch.permute(x, (2,3, 0, 1))

        # f_dash = torch.multiply(x, mc)

        # f_dash.permute((2, 3, 0, 1))
        
        return  mc

I want to infer the input for this layer from the previous layer, it is like a LazyLinear approach,

I run this test and it is totally running:

input = torch.randn(32, 1664, 7, 7)

# GAP = nn.AdaptiveAvgPool2d((1,1))

# GAP(input).shape

# torch.squeeze(input, 1).shape

cam = CAM_infer()

mc = cam(input)

if mc.ndim != 2:

raise ValueError('Output dimension should be 2-D')

print(mc.shape)

But when I transfer the data and the model to CUDA side I got this error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

May anyone help? :eyes:

Thanks a lot for your continuous support :heart:

2

it seems there is a problem regarding lazyLinear parts, cause in the forward method now I transfered them to the cuda side before creating an instance of the CAM_infer and it seems it is working properly :thinking:

May any one elaborate?

3

Yes, the problem is expected to be with the LazyLinear layers because you are instantiating them repeatedly each time the forward pass executes, and by default they would have device CPU.

Your other layers are correctly moved to CUDA because they were declared in the __init__ method, and calling .cuda() on the module should move these layers. But as your LazyLinear layers were not already created, they would not be moved.

Note that this also causes another problem—training this model wouldn’t actually train the LazyLinear layers because they are being recreated and reinitialized every forward pass.