How to do weight normalization in last classification layer?

username: wenying_dai
cooked: <p>Many loss functions in face recognition will norm feature or norm weight before calculate softmax loss, such as normface(<a href="https://arxiv.org/abs/1704.06369" rel="nofollow noopener">https://arxiv.org/abs/1704.06369</a>), L2-softmax(<a href="https://arxiv.org/abs/1703.09507" rel="nofollow noopener">https://arxiv.org/abs/1703.09507</a>) and so on.</p>
<p>I’d like to know how to norm weight in the last classification layer.</p>
<pre><code class="lang-auto">self.feature =  torch.nn.Linear(7*7*64, 2) # Feature extract layer
self.pred = torch.nn.Linear(2, 10, bias=False) # Classification layer
</code></pre>
<p>I want to replace the weight parameter in self.pred module with a normalized one.<br>
In another word, I want to replace weight in-place, like this:</p>
<pre><code class="lang-auto">self.pred.weight = self.pred.weight / torch.norm(self.pred.weight, dim=1, keepdim=True)
</code></pre>
<p>When I trying to do this, there is something wrong:</p>
<pre><code class="lang-auto">TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)
</code></pre>
<p>I am new comer to pytorch, I don’t know what is the standard way to handle this. Thanks a lot!</p>
<p>Here is the whole code:</p>
<pre><code class="lang-auto">class Model(torch.nn.Module):

    def __init__(self):
        super(Model, self).__init__()

        self.backbone = torch.nn.Sequential(

            torch.nn.Conv2d(1, 8, kernel_size=7, stride=1, padding=3),

            # 28 * 28
            torch.nn.Conv2d(8, 16, kernel_size=5, stride=1, padding=2),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2, kernel_size=2),

            # 14 * 14
            torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(stride=2, kernel_size=2),

            # 7 * 7
            torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),

        )

        self.feature =  torch.nn.Linear(7*7*64, 2) # Feature extract layer
        self.pred = torch.nn.Linear(2, 10, bias=False) # Classification layer

        ## something wrong here
        self.pred.weight = self.pred.weight / torch.norm(self.pred.weight, dim=1, keepdim=True)


        for m in self.modules():
            if isinstance(m, torch.nn.Conv2d):
                torch.nn.init.xavier_normal_(m.weight)
            elif isinstance(m, torch.nn.BatchNorm2d):
                torch.nn.init.constant_(m.weight, 1)
                torch.nn.init.constant_(m.bias, 0)


    def forward(self, x):

        x = self.backbone(x)
        x = x.view(-1, 7 * 7 * 64)
        x = self.feature(x)
        x = self.pred(x)

        return x
</code></pre>
<p>Besides,</p>
<ol>
<li>if I don’t want to use <code>view</code> function in forward function, how to deal with it in <code>__init__</code> function？</li>
<li>how I can extract the weight, which is before normalization and after normalization?</li>
<li>If I want to norm feature which is the output feature map in the last two fc layer  in <code>__init__</code> function, how to do it?</li>
</ol>
<p>Thanks a lot!</p>

username: wenying_dai
cooked: <p>Besides, I have read the doc about <code>WeightNorm</code>(<a href="https://pytorch.org/docs/stable/_modules/torch/nn/utils/weight_norm.html" rel="nofollow noopener">https://pytorch.org/docs/stable/_modules/torch/nn/utils/weight_norm.html</a>)</p>
<p>I just need one normalized weight, and do not need both of <code>g</code> and <code>v</code>.</p>

username: ptrblck
cooked: <p>This code should work:</p>
<pre><code class="lang-python">        self.pred = torch.nn.Linear(2, 10, bias=False)

        with torch.no_grad():
            self.pred.weight.div_(torch.norm(self.pred.weight, dim=1, keepdim=True))
        ...
</code></pre>
<ol>
<li>You have to flatten the activation somehow, so <code>.view</code> would be the easiest way.<br>
Alternatively, you could write a <code>Flatten</code> module, initialize if in your model’s <code>__init__</code>, and call it in your <code>forward</code>. I’m not sure, if I understood your question correctly, so let me know if I missed something.</li>
<li>You can just access it as you’ve already done normalizing the weights: <code>print(self.pred.weight)</code>.</li>
<li>You could norm the activation in the <code>forward</code> method (similar to your weight norm code).</li>
</ol>

username: wenying_dai
cooked: <p>I am still not clear about " with torch.no_grad():"</p>
<pre><code class="lang-auto">with torch.no_grad():
            self.pred.weight.div_(torch.norm(self.pred.weight, dim=1, keepdim=True))
</code></pre>
<p>Is this means that the parameter of weight will not be update?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/b/bf747804e50e2789d7bfff33f0447342e8b86840.png" data-download-href="https://discuss.pytorch.org/uploads/default/bf747804e50e2789d7bfff33f0447342e8b86840" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/b/bf747804e50e2789d7bfff33f0447342e8b86840_2_667x499.png" alt="image" data-base62-sha1="rjGJgK0v1Hr2t9ulIho6YS6niFy" width="667" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/b/bf747804e50e2789d7bfff33f0447342e8b86840_2_667x499.png, https://discuss.pytorch.org/uploads/default/optimized/2X/b/bf747804e50e2789d7bfff33f0447342e8b86840_2_1000x748.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/b/bf747804e50e2789d7bfff33f0447342e8b86840.png 2x" data-dominant-color="F3F7FA"><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">1060×794 30.4 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>
<p>Acutally, this is what I want. And I want to know how to get</p>
<ol>
<li>param: weight which is not normalized</li>
<li>the weight which is used in fc calculation</li>
</ol>

username: ptrblck
cooked: <p>The <code>torch.no_grad()</code> guard just makes sure that the operations in this block won’t be recorded by Autograd. The parameter will still be updated in your main training loop.</p>
<p>It sounds like points 1. and 2. are referring to the same parameters.<br>
You can get the weight used in the linear layer with:</p>
<pre><code class="lang-python">self.pred,weight  # inside the model class
# or
model.pred.weight  # outside the model class
</code></pre>

username: sharifza
cooked: <p>Don’t you need to apply that in every forward step? here it seems like you just normalize it after initialization but the gradients might change it…</p>

username: jfyeh
cooked: <p>Hi ptrblck,</p>
<p>I also need to do weight initialization in current project, but I need to do it in every forward pass.<br>
So the torch.no_grad() method is not suit for me.</p>
<p>I found the solution in <a href="https://github.com/pytorch/pytorch/issues/16207#issuecomment-456058059" rel="nofollow noopener">here</a>.</p>
<pre><code class="lang-auto">self.pred.weight = torch.nn.Parameter(self.pred.weight / torch.norm(self.pred.weight, dim=1, keepdim=True))
</code></pre>
<p>I wanna know those cast operation(cast Parameter to Tensor) will affect the gradient flow or not ?<br>
Or there is a better way to do weight normalization in forward pass ?</p>
<p>Thank you so much.</p>

username: ptrblck
cooked: <p>I’m not sure what your workflow is, as you are initializing the parameter in each forward pass, thus removing whatever was changed during training.</p>
<p>If you don’t need to update this parameter, you might just use a plain tensor for its operations.</p>

username: jfyeh
cooked: <p>Hi, thank you for your quick reply.</p>
<p>I am using the loss functions borrowed from face recognition modules(like <a href="https://arxiv.org/abs/1801.09414" rel="nofollow noopener">CosFace</a>, <a href="https://arxiv.org/pdf/1801.07698.pdf" rel="nofollow noopener">ArcFace</a>, etc).</p>
<p>In those loss function, they will normalize both weights of the last linear layer and features before the matmul operation. (for calculate the cosine values)<br>
If I understand this correctly, I am not trying to initialize those weights, I need both of weights and features are normalized before the matmul operation.<br>
So I think the normalize operation should also be considered when calculating the gradient ?<br>
That’s why I said torch.no_grad() is not suit for me.</p>
<p>Please correct me if I misunderstood something, many thanks.</p>

username: Ashima_Garg
cooked: <p>Hi, I tried this with bias set to default value i.e. True.<br>
However, I got this error:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation. 
</code></pre>
<p>When I removed <code>with torch.no_grad()</code>, I got this error:</p>
<pre><code class="lang-auto"> RuntimeError: a leaf Variable that requires grad is being used in an in-place operation. 
</code></pre>
<p>Any ideas on what else can be tried out?</p>

username: ptrblck
cooked: <p>Could you post a code snippet to reproduce the error, please?</p>

username: Ashima_Garg
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, Thanks for your quick response!</p>
<p>Following is the code snippet that reproduces the error.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim


class LinearLayer(nn.Module):
    def __init__(self):
        super(LinearLayer, self).__init__()
        self.classifier = nn.Sequential(nn.Linear(10, 6))

    def forward(self, x):
        with torch.no_grad():                     ### 1
            self.classifier[0].weight.div_(torch.norm(self.classifier[0].weight, dim=1, keepdim=True))
        out = self.classifier(x)
        with torch.no_grad():                     ### 2
            self.classifier[0].weight.div_(torch.norm(self.classifier[0].weight, dim=1, keepdim=True))
        return out

linear = LinearLayer()
optimizer = optim.SGD([{'params': linear.classifier.parameters(), 'lr': 0.1},], 
                      momentum=0.9, weight_decay=5e-4)


for i in range(10):
    x = torch.randn(4, 10)
    targets = torch.LongTensor([1, 2, 3, 4])
    optimizer.zero_grad()
    out = linear(x)
    print(out.size())
    print(targets.size())
    loss = nn.CrossEntropyLoss()(out, targets)
    loss.backward()
    optimizer.step()    
</code></pre>
<p>Following is the error that this code will give.</p>
<pre><code class="lang-auto">Error: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [10, 6]], which is output 0 of TBackward, is at version 3; expected version 2 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
</code></pre>
<p>For my use case, I need to normalize the weights of the classification layer before (<code>### 1</code>) and after (<code>### 2</code>) forward pass. I need to do few more operations on the weights after the <code>###2</code>.</p>
<p>The code will work fine if I remove the second normalization i.e. <code>###2</code>.</p>
<p>Before I looked at this post, I had posted few alternatives for this normalization <a href="https://discuss.pytorch.org/t/l2-normalizing-the-weights/141006">here</a>. Please can you verify if my Implementation 1 is correct in that post?</p>

username: ptrblck
cooked: <p>I don’t think <code>### 2</code> is a valid operation as you are manipulating the weigths inplace after using them, which would update the weights and make the forward activations stale. Could you create a new parameter after <code>### 2</code> which is used in the needed operations?</p>

username: Ashima_Garg
cooked: <p>Hi, Thanks again!<br>
Will this be the right way to do that? Here I have created variable for both operations <code>###1</code> and <code>###2</code></p>
<pre><code class="lang-auto">def __init__(self):
    super(MyModel, self).__init__()
    self.linear = nn.Linear(2, 2)

def forward(self, x):
    weight = F.normalize(self.linear.weight)                  ### 1
    out = torch.mm(x, weight.t()) + self.linear.bias
    weight = F.normalize(weight)                              ### 2
    return out
</code></pre>
<p>Will this be the same as that of first <code>###1</code> in place operation followed by creating a new parameter for <code>###2</code> operation.</p>
<pre><code class="lang-auto">def __init__(self):
    super(MyModel, self).__init__()
    self.linear = nn.Linear(2, 2)

def forward(self, x):
    with torch.no_grad():                     ### 1
            self.classifier[0].weight.div_(torch.norm(self.classifier[0].weight, dim=1, keepdim=True))
    out = torch.mm(x, weight.t()) + self.linear.bias
    weight = F.normalize(weight)             ### 2
    return out
</code></pre>
<p>Thanks,</p>

username: ptrblck
cooked: <p>I guess both approaches should work (you should get an error message if they are causing other issues), but they won’t be the same. In the second approach you are manipulating the <code>weight</code> of <code>self.classifier[0]</code> inplace, which would also change its execution in case you are calling <code>self.classifier[0](input)</code> at one point and would also be visible in the <code>state_dict</code> if you store the model.</p>

username: Vishhvak_Srinivasan
cooked: <p>Did you figure out how to implement weight normalization for acrface? Would really like to know, since I’m struggling with the same issue.</p>