Nested custom module with nn parameter not showing increase in parameters

username: Salvi
cooked: <p>Hi,</p>
<p>I made a custom module (noise injector) as shown below:</p>
<pre><code class="lang-auto">class NoiseLayer(nn.Module):


    def __init__(self, shape):
        super(NoiseLayer, self).__init__()
        
        self.nw = nn.Parameter(torch.randn(shape, device=device)) #move tensor to device before parameter call
        self.n = torch.randn(shape, device=device).normal_(mean=0,std=1)

        
    def forward(self, x):

        return x + self.nw*self.n
</code></pre>
<p>I aim to apply the noise layer per convolution in a module list found in another custom module. However, when I check the number of trainable parameters, it does not increase.</p>
<pre><code class="lang-auto">    def forward(self, x):

        for level, conv in enumerate(nn.ModuleList(etc,etc,etc)):
            x = conv(x)
            self.a = nn.ModuleList([NoiseLayer(x.size())])
            x = self.a[0](x)
            x = self.pooling[level](x)
</code></pre>
<p>I am not sure what I am doing wrong. One approach I also tried is placing the ‘shape’ variable to ‘forward’ and moving self.a to above. This approach increased my forward/backward pass size but it did not change my total params.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Salvi" data-post="1" data-topic="207419">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/salvi/48/65749_2.png" class="avatar"> Salvi:</div>
<blockquote>
<p>However, when I check the number of trainable parameters, it does not increase.</p>
</blockquote>
</aside>
<p>This is expected, since <code>self.n</code> is not a trainable parameter.<br>
Either register it via <code>self.n = nn.Parameter(...)</code> if you want to train it or via <code>self.register_buffer</code> if you do not want to train it, but still register it (it will then show up in the <code>state_dict</code>).</p>

username: Salvi
cooked: <p>Hi,</p>
<p>Thanks for the reply! Somehow, it still doesn’t work…</p>
<pre><code class="lang-auto">class NoiseLayer(nn.Module):


    def __init__(self, shape, noise_std):
        super(NoiseLayer, self).__init__()
        
        self.nw = nn.Parameter(torch.randn(shape, device=device)) #move tensor to device before parameter call
        self.n = nn.Parameter(torch.randn(shape, device=device).normal_(mean=0,std=noise_std))

    def forward(self, x):

        return x + self.nw*self.n

</code></pre>
<p>I tried placing it in a module list too but that also didn’t work</p>
<pre><code class="lang-auto">        self.noiselayers = NoiseLayer
        self.noise = noise
    def forward(self, x):

        for level, conv in enumerate(self.encoder):
            x = conv(x)
            if self.noise:
              x = self.noiselayers(x.size(), self.noise)(x)
            x = self.pooling[level](x)

</code></pre>

username: ptrblck
cooked: <p>It works for me:</p>
<pre data-code-wrap="python"><code class="lang-python">class NoiseLayer(nn.Module):
    def __init__(self, shape, noise_std):
        super(NoiseLayer, self).__init__()
        device = "cpu"
        self.nw = nn.Parameter(torch.randn(shape, device=device)) #move tensor to device before parameter call
        self.n = nn.Parameter(torch.randn(shape, device=device).normal_(mean=0,std=noise_std))

    def forward(self, x):

        return x + self.nw*self.n
    
layer = NoiseLayer((1, 1), 1)
for name, param in layer.named_parameters():
    print(name, param)
# nw Parameter containing:
# tensor([[2.0764]], requires_grad=True)
# n Parameter containing:
# tensor([[0.0912]], requires_grad=True)

print(len(list(layer.parameters())))
# 2
</code></pre>

username: Salvi
cooked: <p>It seems to work standalone but when I place it inside a module it doesn’t seem to register</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.nn import functional as F
import numpy as np
from torchsummary import summary

if torch.cuda.is_available():
  device = torch.device("cuda")
  print(torch.cuda.get_device_name(0))
else:
  device = torch.device("cpu")
  print('no gpu available')


class ConvBlock(nn.Module):

    def __init__(self, ndims, in_channels, out_channels, ks, stride=1, p=2):
        super().__init__()

        Conv = getattr(nn, 'Conv%dd' % ndims)
        self.main = Conv(in_channels, out_channels, ks, stride, p)
        self.activation = nn.LeakyReLU(0.2)

    def forward(self, x):
        out = self.main(x)
        out = self.activation(out)
        return out


class NoiseLayer(nn.Module):


    def __init__(self, shape, noise_std):
        super(NoiseLayer, self).__init__()
        
        self.nw = nn.Parameter(torch.randn(shape, device=device)) #move tensor to device before parameter call
        self.n = nn.Parameter(torch.randn(shape, device=device).normal_(mean=0,std=noise_std))

    def forward(self, x):

        return x + self.nw*self.n

class Dnet(nn.Module):

    def __init__(self,
                 inshape=None,
                 input_feats=None,
                 enc_features=[4,6],
                 max_pool=2, noise=False):

        super().__init__()
        ndims = len(inshape)
        self.encoder = nn.ModuleList()
        prev_feats = input_feats
        ks = 5 
        for i in range(len(enc_features)):
            convs = ConvBlock(ndims, prev_feats, enc_features[i], ks)
            prev_feats = enc_features[i]
            self.encoder.append(convs)


        #pooling
        self.pooling = nn.ModuleList()
        Conv = getattr(nn, 'Conv%dd' % ndims)
        scale = 2
        for i in range(len(enc_features)):
            pools = Conv(enc_features[i], enc_features[i], ndims, scale, 1)
            self.pooling.append(pools)

        self.noiselayers = NoiseLayer
        self.noise = noise
    def forward(self, x):

        # encoder forward pass
        first_x = x
        x_history = [x]
        for level, conv in enumerate(self.encoder):
            x = conv(x)
            if self.noise:
                    x = self.noiselayers(x.size(), self.noise)(x)
            x_history.append(x)
            x = self.pooling[level](x)


        return x

model = Dnet(inshape=[224,128,128],enc_features=[2,4,8],input_feats=1, noise=1e-1).to(device)
summary(model, input_size=[(1,128,128,128)])

</code></pre>
<p>If I set noise=False and noise=1, it will still output the same amount of parameters (or is it because I’m using torchsummary?)</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/a/bad6d8ae29bc817b2375c2fbb5d98e72eb374465.png" data-download-href="https://discuss.pytorch.org/uploads/default/bad6d8ae29bc817b2375c2fbb5d98e72eb374465" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/a/bad6d8ae29bc817b2375c2fbb5d98e72eb374465.png" alt="image" data-base62-sha1="qER7basxLdsbEwJuQVirHLp7aPH" width="598" height="500" data-dominant-color="414141"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">693×579 13.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Salvi" data-post="5" data-topic="207419">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/salvi/48/65749_2.png" class="avatar"> Salvi:</div>
<blockquote>
<p>but when I place it inside a module it doesn’t seem to register</p>
</blockquote>
</aside>
<p>You are not registering it properly, since you are initializing it on-the-fly.<br>
Initialize the layer in the <code>__init__</code> to create all parameters.</p>

username: Salvi
cooked: <p>Yes, so I did not want to do that as x.size() changes and placing on the init meant that I have to calculate each size manually (or should I bite the bullet and do this each time?)</p>
<p>Alternatively, I did some digging and found something like this would make it register (increase forward pass size) but not add trainable parameters:</p>
<pre><code class="lang-auto">class NoiseLayer(nn.Module): 

    def __init__(self):
        super().__init__()
        self.register_parameter('nw', nn.Parameter(torch.randn(1)))
        self.register_parameter('n', nn.Parameter(torch.randn(1)))
        

    def forward(self, x, shape, noise_std):  
        if self.n is nn.Parameter(torch.randn(1)):
            self.n = torch.randn(shape).normal_(mean=0, std=noise_std)
        if self.nw is nn.Parameter(torch.randn(1)):
            self.nw = torch.nn.Parameter(torch.randn(shape))
        return x + self.nw*self.n
</code></pre>
<p>I initialized this one*</p>
<pre><code class="lang-auto">        self.noiselayers = nn.ModuleList()
        for i in range(len(enc_features)):
            self.noiselay = NoiseLayer()
            self.noiselayers.append(self.noiselay)
</code></pre>