Optimizer.step() disregards learning rate with multiple nn.Parameter()

username: amirhf
cooked: <p>I have a custom module with two sets of nn.Parameters():</p>
<pre><code class="lang-auto">class CustomLayer(nn.Module):
    def __init__(self):
        self.param1 = nn.Parameter(torch.rand())
        self.param2 = nn.Parameter(torch.rand())
</code></pre>
<p>later when I optimize with optimizer.step():</p>
<pre><code class="lang-auto">layer = CustomLayer()
optimizer = torch.optim.SGD(layer.parameters(), lr=0.5)
optimizer.step()
</code></pre>
<p>it doesn’t do <code>param = param - 0.5*param.grad</code> it instead calculates: <code>param = param - param.grad</code></p>

username: albanD
cooked: <p>Hi,</p>
<p>Can you give a code sample that reproduces this issue please?</p>

username: amirhf
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>, yes here is a sample code. I’m trying to implement separate loss in <a href="http://proceedings.mlr.press/v101/li19b.html" rel="nofollow noopener">separate loss paper</a></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
torch.manual_seed(3275)


class SeparateLoss(nn.Module):
    def __init__(self, num_classes, feat_dim):
        super(SeparateLoss, self).__init__()
        self.centers = torch.randn(num_classes, feat_dim)
        self.intra_centers = nn.Parameter(self.centers.view_as(self.centers))
        self.inter_centers = nn.Parameter(self.centers.view_as(self.centers))
        self.cosine_similarity = torch.nn.CosineEmbeddingLoss()
        self.reset_params()

    def reset_params(self):
        nn.init.kaiming_normal_(self.centers.data)

    def forward(self, f, label):
        batch_size = label.size(0)
        y = label.new_ones(1, batch_size).float()
        s = self.centers.size(0)

        # calculating intra loss
        # ----------------------
        # fix the scales for c_intra gradients as in eq (6)
        counts = self.intra_centers.new_ones(s)
        ones = self.intra_centers.new_ones(batch_size)
        counts.scatter_add_(0, label.long(), ones)
        intra_hook = self.intra_centers.register_hook(lambda grad: grad/counts.unsqueeze(1))

        # eq (2)
        batch_centers = self.intra_centers.index_select(0, label.long())
        Lintra = (-1/2) * (-1 * self.cosine_similarity(f, batch_centers, y) + 2) + 1

        # calculating inter loss
        # ----------------------
        # fix the scales for c_inter gradients as in eq (7)
        inter_hook = self.inter_centers.register_hook(lambda grad: grad/2.0)
        dist = torch.nn.functional.cosine_similarity(
                self.inter_centers.unsqueeze(1).expand(s, s, self.inter_centers.size(1)),
                self.inter_centers.unsqueeze(0).expand(s, s, self.inter_centers.size(1)),
                dim=2
            )

        Linter = (((dist + 1).sum()) - (2.0 * s))/(2*s*(s - 1))

        # calculating separate loss
        Lsep = Lintra + Linter

        return Lsep, intra_hook, inter_hook


def unit_test():

    device = torch.device('cpu')
    m = 8  # number of samples
    k = 3  # number of classes
    feat_dim = 4

    x = torch.randn(m, feat_dim, requires_grad=True).to(device)
    target = torch.tensor([0, 0, 1, 0, 2, 2, 1, 0]).to(device)

    criterion = SeparateLoss(num_classes=k, feat_dim=feat_dim).to(device)
    center_optimizer = torch.optim.SGD(criterion.parameters(), lr=0.5)
    feature_optimizer = torch.optim.SGD({x}, lr=0.1)

    # epoch 1
    # -------
    lsep, intra_hook, inter_hook = criterion(x, target)
    print(lsep.item())

    # backward
    center_optimizer.zero_grad()
    lsep.backward()

    # add intra and inter gradients
    print(f'x.grad = {x.grad}')
    centers_grad = 0
    for param in criterion.parameters():
        centers_grad += param.grad
    # update centers grad
    for param in criterion.parameters():
        param.grad = centers_grad
        print(f'c.grad = {param.grad}')

    # update parameters
    feature_optimizer.step()
    center_optimizer.step()

    # release hooks
    intra_hook.remove()
    inter_hook.remove()


if __name__ == "__main__":
    unit_test()

</code></pre>
<p>so <code>center_optimizer.step()</code> should do <code>param = param - 0.5*param.grad</code> but when I debug the calcualted update is <code>param = param - param.grad</code></p>

username: albanD
cooked: <p>You should replace the following:<br>
(you should not give the same Tensor to different parameters. otherwise they have the same memory and so will always have the same value. And if you do an update on one then the other, then you do 2x the update)</p>
<pre><code class="lang-auto">        self.intra_centers = nn.Parameter(self.centers.view_as(self.centers))
</code></pre>
<p>by</p>
<pre><code class="lang-auto">        self.intra_centers = nn.Parameter(self.centers.clone())
</code></pre>
<p>And (you should not use .data)</p>
<pre><code class="lang-auto">        nn.init.kaiming_normal_(self.centers.data)
</code></pre>
<p>by</p>
<pre><code class="lang-auto">        nn.init.kaiming_normal_(self.centers)
</code></pre>

username: amirhf
cooked: <p>Thank you <a class="mention" href="/u/alband">@albanD</a>. This worked for me. Everything is fine until first epoch. I have problem in the second epoch though.<br>
Following the question I asked yesterday and you answered:</p>
<p>my model has two sets of parameters <code>a</code> and <code>b</code> which require gradients. I have two different loss functions as:</p>
<pre><code class="lang-auto">loss1 = function(a, b)
loss2 = function(b)
</code></pre>
<p>and the total loss is</p>
<pre><code class="lang-auto">total_loss = loss1 + loss2
</code></pre>
<p>Assuming:</p>
<ul>
<li>
<code>loss1</code> calculates <code>b1.grad</code>
</li>
<li>
<code>loss2</code> calculates <code>b2.grad</code>,<br>
then the total <code>b.grad</code> calculated by <code>total_loss.backward()</code> is <code>b.grad = b1.grad + b2.grad</code>.</li>
</ul>
<p>My goal is to modify <code>b1.grad</code> coming from <code>loss1</code> and <code>b2.grad</code> comig from <code>loss2</code> and then add them together as backward gradients for <code>b</code>. Currently when I do <code>total_loss.bakward()</code>, it already gives me the accumulated gradients for <code>b.grad=b1.grad + b2.grad</code>. How can I access and modify each individual <code>b1.grad</code> and <code>b2.grad</code> so the <code>total_loss.backward()</code> returns the modified and added <code>b1.grad</code> and <code>b2.grad</code>?<br>
You told me to make intermediate gradients which I did in this code sample. However, I’m not sure if</p>
<pre><code class="lang-auto">centers_grad = 0
    for param in criterion.parameters():
        centers_grad += param.grad
    # update centers grad
    for param in criterion.parameters():
        param.grad = centers_grad
        print(f'c.grad = {param.grad}')
</code></pre>
<p>is doing the job of <code>b.grad = b1.grad + b2.grad</code>. It is ok for first epoch but doing the same thing in second epoch already adds the gradients for me. It seems like the <code>centers_grad</code> variable gets attached to something. Do you have any idea how to fix this?</p>

username: albanD
cooked: <p>Given that you modify <code>.grad</code> with things that <code>requires_grad=True</code>, I guess the next backward keep the graph from the previous iterations? Can you try setting all the <code>.grad</code> to None at the end of your iteration (instead of calling <code>.zero_grad()</code>.</p>