How to create a custom layer that accepts the input during forward pass

theairbend3r: <p>For instance, if I have a network as follows -</p>
<pre><code class="lang-auto">class SomeClassifier(nn.Module):
    def __init__(self):
        super(SomeClassifier, self).__init__()

        # all three conv blocks maintain the original image size.
        self.cnn1 = nn.Conv2d()
        self.cnn2 = nn.Conv2d()
        self.cnn3 = nn.Conv2d()
        self.cnn4 = nn.Conv2d()
        self.fc = nn.Linear()

    def forward(self, img):
        img_rep = self.cnn1(img)
        x = self.cnn2(img_rep)
        y = self.cnn3(img_rep)
        z = self.cnn4(img_rep)

        # perform custom operation on x, y, and z and merge them into a single tensor called p. p is the same shape as x, y, and z.

        p = torch.concat(x, y, z) # or average the 3 tensors so that we have a new tensor that is the same shape as x,y, and z.
        p = torch.flatten()
        p = self.fc(p)
        
        return p
</code></pre>
<p>The point being, If I have an input image:</p>
<ul>
<li>I pass it through 3 different conv filters.</li>
<li>I perform some custom operations on the output of each of the 3 conv filters.</li>
<li>Finally, I merge the resulting 3 feature maps into a single tensor. The shape of this new filter is the same size as the <code>img_rep</code>.</li>
</ul>
<p>This is clearly a block of operation I can perform again and again. I cannot use <code>nn.Sequential</code> for this because I have a multiple  outputs here. I need a function. So, can I do the following?</p>
<pre><code class="lang-auto">class SomeClassifier(nn.Module):
    def __init__(self):
        super(SomeClassifier, self).__init__()

        # all three conv blocks maintain the original image size.
        self.cnn1 = nn.Conv2d()
        self.fc = nn.Linear()

    def forward(self, img):
        img_rep = self.cnn1(img)
        p = self.custom_module(input_img_rep=img_rep, in_channel, out_channel)
        p = torch.flatten()
        p = self.fc(p)
        
        return p



    def custom_module(self, input_img_rep, input_channels_to_this_block, output_channels_from_this_block):
        cnn2 = nn.Conv2d()
        cnn3 = nn.Conv2d()
        cnn4 = nn.Conv2d()

        x = cnn2(img_rep)
        y = cnn3(img_rep)
        z = cnn4(img_rep)

        # perform custom operation on x, y, and z and merge them into a single tensor called p. p is the same shape as x, y, and z.

        o = torch.concat(x, y, z) # or average the 3 tensors so that we have a new tensor that is the same shape as x,y, and z.
        
        return o
</code></pre>
<p>When I do this, I have issues with the GPU. Basically, the weights for the <code>custom_block</code> are still on the CPU while the weights of the rest of the network are on the GPU.</p>
<p>This is because I’m not initializing this <code>custom_block</code> method in the <code>__init__()</code>. So, let’s say I do that.</p>
<pre><code class="lang-auto">class SomeClassifier(nn.Module):
    def __init__(self):
        super(SomeClassifier, self).__init__()

        # all three conv blocks maintain the original image size.
        self.cnn1 = nn.Conv2d()
        self.custom_block = self.custom_module(...)
        self.fc = nn.Linear()

    def forward(self, img):
        img_rep = self.cnn1(img)
        p = self.custom_block(...)
        p = torch.flatten()
        p = self.fc(p)
        
        return p



    def custom_module(self, input_img_rep, input_channels_to_this_block, output_channels_from_this_block):
        cnn2 = nn.Conv2d()
        cnn3 = nn.Conv2d()
        cnn4 = nn.Conv2d()

        x = cnn2(img_rep)
        y = cnn3(img_rep)
        z = cnn4(img_rep)

        # perform custom operation on x, y, and z and merge them into a single tensor called p. p is the same shape as x, y, and z.

        o = torch.concat(x, y, z) # or average the 3 tensors so that we have a new tensor that is the same shape as x,y, and z.
        
        return o
</code></pre>
<p>The problem with this is, my <code>custom_block</code> requires an input image representation. So, I can’t really initialize it in <code>__init__()</code>. I have to directly use in the <code>forward()</code> method (which leads to not all weights being transferred to the GPU.)</p>
<p>So, how do I go about using this <code>custom_block()</code> in PyTorch?</p>

albanD: <p>Hi,</p>
<p>You should make your custom module an actuall independent nn.Module. And that will solve your issues <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
Instantiate it in the <code>__init__</code>  and use it in the forward.</p>

theairbend3r: <p>Ah. Okay. How do we implement that? Something like the following?</p>
<p>Custom module.</p>
<pre><code class="lang-auto">class CustomModule(nn.Module):
    def __init__(self, input_channels_to_this_block, output_channels_from_this_block):
        super(CustomModule, self).__init__()

        self.cnn2 = nn.Conv2d()
        self.cnn3 = nn.Conv2d()
        self.cnn4 = nn.Conv2d()



    def forward(self, img_rep):
        x = cnn2(img_rep)
        y = cnn3(img_rep)
        z = cnn4(img_rep)

        # perform custom operation on x, y, and z and merge them into a single tensor called p. p is the same shape as x, y, and z.

        o = torch.concat(x, y, z) # or average the 3 tensors so that we have a new tensor that is the same shape as x,y, and z.
        
        return o
</code></pre>
<p>Classifier which uses the custom module.</p>
<pre><code class="lang-auto">class SomeClassifier(nn.Module):
    def __init__(self):
        super(SomeClassifier, self).__init__()

        # all three conv blocks maintain the original image size.
        self.cnn1 = nn.Conv2d()
        self.custom_block = CustomModule(input_channels_to_this_block, output_channels_from_this_block)
        self.fc = nn.Linear()

    def forward(self, img):
        img_rep = self.cnn1(img)
        p = self.custom_block(img)
        p = torch.flatten()
        p = self.fc(p)
        
        return p
</code></pre>

albanD: <p>Yes exactly !<br>
That way you can reuse your custom block wherever you need it.</p>