Weight norm calculation

username: vainaijr
cooked: <p>How is weight normalization calculated?</p>
<pre><code class="lang-auto">import torch, torch.nn as nn
lin = nn.Linear(3, 3, bias=False)
inp = torch.randn(3, 3)
lin = nn.utils.weight_norm(lin)
optimizer = torch.optim.SGD(lin.parameters(), lr=0.01)
list(lin.parameters())
</code></pre>
<p>[Parameter containing:<br>
tensor([[0.6400],<br>
[0.6961],<br>
[0.6579]], requires_grad=True), Parameter containing:<br>
tensor([[ 0.2672,  0.5720,  0.1048],<br>
[-0.4538,  0.5243, -0.0612],<br>
[-0.5537,  0.0813, -0.3459]], requires_grad=True)]</p>
<pre><code class="lang-auto">optimizer.zero_grad()
loss = lin(inp).sum()
loss.backward()
optimizer.step()
</code></pre>
<pre><code class="lang-auto">lin.weight_g
</code></pre>
<p>Parameter containing:<br>
tensor([[0.6161],<br>
[0.6924],<br>
[0.6708]], requires_grad=True)</p>
<pre><code class="lang-auto">lin.weight_v
</code></pre>
<p>Parameter containing:<br>
tensor([[ 0.2619,  0.5746,  0.1042],<br>
[-0.4715,  0.5084, -0.0660],<br>
[-0.5581,  0.0611, -0.3437]], requires_grad=True)</p>
<pre><code class="lang-auto">lin.weight
</code></pre>
<p>tensor([[ 0.2672,  0.5720,  0.1048],<br>
[-0.4538,  0.5243, -0.0612],<br>
[-0.5537,  0.0813, -0.3459]], grad_fn=)</p>
<p>I use this formula to calculate lin.weight from lin.weight_g and lin.weight_v, but it gives wrong lin.weight</p>
<pre><code class="lang-auto">lin.weight = (lin.weight_g/lin.weight_v.norm())*lin.weight_v
</code></pre>

username: albanD
cooked: <p>If you check the <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.utils.weight_norm">doc</a>, the default for <code>dim</code> is <code>0</code>, maybe you want to set it to <code>None</code> here?</p>

username: vainaijr
cooked: <p>Yes, it works for dim=None, in weight_norm, also, for default dim=0,<br>
I used this formula,</p>
<pre><code class="lang-auto">lin.weight_g*(lin.weight_v/lin.weight_v.norm(dim=1, keepdim=True))
</code></pre>
<p>or even this works</p>
<pre><code class="lang-auto">lin.weight_g*(lin.weight_v/torch.norm_except_dim(lin.weight_v, 2, dim=0))
</code></pre>
<p>I found out that incase of gradient descent here, weight_g, weight_v get updated, and after optimizer.step(), weight stays as it is. (so maybe, last training batch will have no impact on weights?)<br>
weight is calculated from these updated weight_g, weight_v, before forward.</p>
<p>also, the first lin.weight_g is calculated as,</p>
<pre><code class="lang-auto">torch.norm_except_dim(lin.weight, 2, dim=0)
</code></pre>
<p>which leads to an output, whose shape is not changed for the specified dimension.<br>
also, the first lin.weight_v value is same as lin.weight value.</p>
<p>I think main thing here is,</p>
<ul>
<li>
<p>torch.norm(dim=1) gives output, whose shape is changed only in the specified dimension, and if we keepdim=True, then shape of that dimension becomes 1.</p>
</li>
<li>
<p>torch.norm_except_dim(dim=1) gives output, whose shape is not changed for the specified dimension, and rest all dimensions shape becomes 1.</p>
</li>
<li>
<p>and we want lin.weight_g, lin.weight_v.norm shape to be same, therefore, we use, lin.weight_v.norm(dim=1, keepdim=True) or torch.norm_except_dim(lin.weight_v, 2, dim=0)</p>
</li>
<li>
<p>if I set dim=1 in weight_norm, then lin.weight_g would be calculated by torch.norm_except_dim(lin.weight, 2, dim=1), and lin.weight would consider lin.weight_v.norm(dim=0, keepdim=True) or torch.norm_except_dim(lin.weight_v, 2, dim=1) to match with shape of lin.weight_g</p>
</li>
</ul>

username: Ayush_Kumar_Agarwal
cooked: <p>Original Weights can simply be calculated by removing the weight norm from the currrent weight norm layer, using the function “torch.nn.utils.remove_weight_norm(layer)”</p>
<p>Please run the below  code snippet once and check output values :</p>
<pre><code class="lang-auto">l = nn.Linear(2, 4)
print(l.state_dict())
w1 = weight_norm(l, name="weight")
print(w1.state_dict())
remove_weight_norm(w1)
print(w1.state_dict())
</code></pre>