Weights are not updated for specific layers

username: hs271095
cooked: <p>Hello I am new to pytorch and trying to implement a simple model with an attention layer. On inspection I am able to see that my model is not updating the weights for the attention layers in the model(code below). I have plotted the gradient flow and also used model.parameters() to check for weights before and after the .backward() and optimizer.step() calls and the weights remain the same.</p>
<p>My model</p>
<pre><code class="lang-auto">class ModelA(nn.Module):
    def __init__(self):
        super(ModelA,self).__init__()
        self.L = 128
        self.D = 64
        self.K = 1
        self.embedding=nn.Linear(24, self.L)
        self.attention = nn.Sequential(
            nn.Linear(self.L, self.D),
            nn.Tanh(),
            nn.Linear(self.D, self.K),
            nn.InstanceNorm1d(19)
        )
        self.classifier= nn.Linear(self.L*self.K, 1)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        embedding = self.embedding(x)
        A = self.attention(embedding)
        A = torch.transpose(A, 2, 1)
        A = F.softmax(A, dim=1)
        M = torch.bmm(A, embedding)
        out = self.classifier(M.view(x.size(0),-1))
        return self.sigmoid(out)
</code></pre>
<p>model Intialization and hyper-parameters</p>
<pre><code class="lang-auto">modelA=ModelA()
optimizer = torch.optim.SGD(modelA.parameters(), lr=0.0001)
loss_func = nn.MSELoss()
</code></pre>
<p>training code</p>
<pre><code class="lang-auto">for it in range(15):
    modelA.train()
    total=len(train_data_loader)*batch_size
    train_loss = 0.
    for minibatch in train_data_loader:
        X, Y1, Y2  = minibatch
        output = modelA(X)
        optimizer1.zero_grad()
        total_loss=loss_func(output,Y1.view(-1,1))
        train_loss += total_loss.item()
#         a = list(modelA.parameters())[2].clone()
        total_loss.backward()
        plot_grad_flow(list(modelA.named_parameters()))
        optimizer.step()
#         b = list(modelA.parameters())[2].clone()
#         print(torch.equal(a.data, b.data))
    train_loss /= total
    print("EPOCH ",it)
    print('Train : Loss: {:.4f}'.format(train_loss))
</code></pre>
<p>Gradient Flow<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/1/8/1833c755de4ab6c7344da0d61a705b31553f81a0.png" alt="image" data-base62-sha1="3s6osDctcA0eZLZ3uctMQ5tmKqs" width="476" height="397"></p>

username: hs271095
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I have been following your comments to learn more about pytorch any help would be appreciated!. Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: ptrblck
cooked: <p>The computation graph doesn’t seem to be detached, but instead the operations seems to be decrease the gradient magnitude for the attention layers.<br>
The <code>embedding</code> layer (which is a linear layer in fact) also skips the <code>attention</code> module, thus gets larger gradients.</p>
<p>Could you explain your use case a bit as well as the input shapes?<br>
I assume you are feeding a 3-dimensional tensor, since the <code>InstanceNorm</code> layer would raise an error for a 2-dim tensor?</p>

username: hs271095
cooked: <p>Thank You for your reply I am working on a multi-task learning problem and this is a model for one of the task.The input to the model is a 3D tensor . I was using the <code>InstanceNorm</code> Layer since i read  in one of the blogs that softmax might make the gradient magnitude small, so its better to use some normalization before doing softmax. I can add more layers to the embedding layer if that would make the gradients increase?</p>