Weights not updated in a custom module

username: seloufian
cooked: <p>I am facing a weird issue when trying to optimize parameters in a custom module. In particular, the weights of certain parameters are not updated, even if their gradients are computed.</p>
<p>My <code>CustomModule</code> is defined as:</p>
<ul>
<li>
<p>A list of Batch Normalization (BN) layers, in which, only one is chosen for the forward/backward pass. That is, the model will learn different gammas/betas (based on the selected BN layer).</p>
</li>
<li>
<p>A fully-connected layer.</p>
</li>
</ul>
<p>Of course, this structure, defined as such, does not make sense. I just wanted to focus on the occurred issue. Below, the structure of the <code>CustomModule</code>.</p>
<pre><code class="lang-auto">class CustomModule(nn.Module):
    def __init__(self, num_batch_norm):
        super(CustomModule, self).__init__()

        self.batch_norm_list = [nn.BatchNorm1d(10) for _ in range(num_batch_norm)]
        self.fc = nn.Linear(10, 2)

    def forward(self, x, num_bn):
        x = self.batch_norm_list[num_bn](x)
        x = self.fc(x)

        return x
</code></pre>
<p>I create a <code>model</code> (instance of <code>CustomModule</code>) with <code>three</code> BN layers and <code>SGD</code> optimizer. After, I define test variables <code>x</code> and <code>y</code> to perform a single pass through the <code>model</code> (including the backward pass).</p>
<pre><code class="lang-auto">torch.manual_seed(0)

# Create a 'CustomModule' with three BN layers.
model = CustomModule(3)

optimizer = torch.optim.SGD(model.parameters(), lr=1e-2)

# Define test variables: the input (x) and the target (y).
x = torch.rand(5, 10)
y = torch.rand(5, 2)

# Pass 'x' through the 2nd (index=1) BN layer.
out = model(x, 1)

loss = torch.linalg.norm(out - y)

model.zero_grad()
loss.backward()
optimizer.step()
</code></pre>
<p>I check that the gradients for the 2nd BN layer (chosen in the forward pass) are computed.</p>
<pre><code class="lang-auto"># Check the weights' gradients in the 2nd BN layer.
model.batch_norm_list[1].weight.grad
# tensor([ 0.0810,  0.0052,  0.2871,  0.3091,  0.2437,  0.0916,  0.1410,  0.4406,
#         -0.0288,  0.0740])
</code></pre>
<p>However, when I check the weights (gammas) after optimization, I noticed that they are not updated (in batch normalization, gammas (weights) are initialized by <code>one</code>).</p>
<pre><code class="lang-auto"># Check the weights in the 2nd BN layer (after optimization)
model.batch_norm_list[1].weight
# Parameter containing:
# tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)
</code></pre>
<p>Having the initial weights (a tensor of ones) and gradients, I expected the updated weights to be like below (using the <code>SGD</code> optimizer).</p>
<pre><code class="lang-auto"># Expected values
weights = model.batch_norm_list[1].weight
gradients = model.batch_norm_list[1].weight.grad
lr = 1e-2

expected_weights = weights - lr * gradients
# tensor([0.9992, 0.9999, 0.9971, 0.9969, 0.9976, 0.9991, 0.9986, 0.9956, 1.0003,
#         0.9993], grad_fn=&lt;SubBackward0&gt;)
</code></pre>
<p>So, I want to know why this issue encounters and how to solve it.</p>

username: Sangyeon_Kim
cooked: <pre><code class="lang-auto">model.zero_grad()
</code></pre>
<p>I think the line should be modified to <code>optimizer.zero_grad()</code></p>

username: JuanFMontesinos
cooked: <p>You are storing the weights in a conventional list which is not tracked by the optimizer.<br>
You should use a pytorvh list.</p>

username: seloufian
cooked: <p><a class="mention" href="/u/sangyeon_kim">@Sangyeon_Kim</a>: Yes, you are right, thanks for the remark.</p>
<p><a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a>: You mean that I have to convert <code>self.batch_norm_list</code> (which holds the BatchNorm layers) into a PyTorch Tensor. But how can I do this? Using <code>torch.tensor(self.batch_norm_list)</code> doesn’t work due to unrecognized element type (in this case, BatchNorm), below is the error message.</p>
<pre><code class="lang-auto">&lt;ipython-input-370-e40c7345c9d8&gt; in __init__(self, num_batch_norm)
      4 
      5         self.batch_norm_list = [nn.BatchNorm1d(10) for _ in range(num_batch_norm)]
----&gt; 6         self.batch_norm_list = torch.tensor(self.batch_norm_list)
      7 
      8         self.fc = nn.Linear(10, 2)

RuntimeError: Could not infer dtype of BatchNorm1d
</code></pre>

username: JuanFMontesinos
cooked: <p>Sorry I should have been more explicit.<br>
You have to use <a href="https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html" class="inline-onebox" rel="noopener nofollow ugc">ModuleList — PyTorch 1.8.1 documentation</a> which is a list-like module which is properly tracked by the engine. Think that attribute of a nn.Module distinct from a tensor or another module is not tracked by the engine.</p>

username: seloufian
cooked: <p>It works! I didn’t know the <code>ModuleList</code>.<br>
Thank you very much.</p>