How to correctly implement in-place Max Norm constraint?

username: lkc
cooked: <p>I’m implementing the max norm constraint as detailed in this <a href="https://github.com/kevinzakka/pytorch-goodies#max-norm-constraint" rel="nofollow noopener">post</a>. Max Norm would constraint the parameters of a given layer based on the L2 norm of the weights.</p>
<p>It’s my understanding that the operations should be done in-place for memory efficiency. I have the following code.</p>
<pre><code class="lang-auto">def forward(self, x):
    x = x.view(-1, n_inputs)
    max_norm_(layer.weight)
    x = layer(x)
return x


def max_norm_(self, w):
    with torch.no_grad():
        norm = w.norm(2, dim=0, keepdim=True)
        desired = torch.clamp(norm, 0, self._max_norm_val)
        w *= (desired / (self._eps + norm))

</code></pre>
<p>Is this the correct way to implementing it? Will this interfere with the optimization step? I’m using <code>torch.no_grad()</code> since we don’t need to get gradients for this operation.</p>
<p>Thanks</p>

username: tom
cooked: <aside class="quote no-group" data-username="lkc" data-post="1" data-topic="96769">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lkc/48/28917_2.png" class="avatar"> lkc:</div>
<blockquote>
<p>It’s my understanding that the operations should be done in-place for memory efficiency. I have the following code.</p>
</blockquote>
</aside>
<p>It’s likely not memory efficiency that matters here (you’ll have other ops that use more memory and you could do this after the optimizer step when memory is less precious than after the forward) but that you change a parameter that has you want inplace.<br>
The implementation you propose looks like the one you linked, but note that it will scale the weight by norm/(norm+eps) when the norm is within the your bound. This may or may not be desired. If not, it might be better to just clamp away both norm and desired from 0 or so.</p>
<p>Best regards</p>
<p>Thomas</p>

username: lkc
cooked: <blockquote>
<p>It’s likely not memory efficiency that matters here (you’ll have other ops that use more memory and you could do this after the optimizer step when memory is less precious than after the forward) but that you change a parameter that has you want inplace.</p>
</blockquote>
<p>That’s a good point. I don’t need to do that operation in the forward step.</p>
<blockquote>
<p>The implementation you propose looks like the one you linked, but note that it will scale the weight by norm/(norm+eps) when the norm is within the your bound. This may or may not be desired. If not, it might be better to just clamp away both norm and desired from 0 or so.</p>
</blockquote>
<p>Thanks for catching that. You are right, norm/(norm + eps) is indeed not the desired behavior. When you say <em>clamp away from both norm and desired from 0</em>, do you mean  something like <code>w *= torch.clamp(desired / (eps + norm), min=1)</code> where we ensure that when <code>norm</code> is within the bounds, there is no effect?</p>

username: lkc
cooked: <p>Nevermind, I get you mean. Just make sure norm or desired are not close to zero.</p>
<p>Thanks, very helpful comments! Appreciate it!</p>

username: tom
cooked: <p>Just to make it explicit (you noticed, but for the next person looking here), clamping the quotient won’t work, but using</p>
<pre><code class="lang-auto">    with torch.no_grad():
        norm = w.norm(2, dim=0, keepdim=True).clamp(min=self._max_norm_val / 2)
        desired = torch.clamp(norm, max=self._max_norm_val)
        w *= (desired / norm)
</code></pre>
<p>should always work as expected and leave small norms alone.</p>