Weight decay only for weights of nn.Linear and nn.Conv*

veritas: <p>In many of the papers and blogs that I read, for example, the recent <a href="https://arxiv.org/abs/2102.06171" rel="noopener nofollow ugc">NFNet</a> paper, the authors emphasize the importance of only including the convolution &amp; linear layer weights in weight decay. Bias values for all layers, as well as the weight and bias values of normalization layers, e.g., LayerNorm, should be excluded from weight decay.<br>
However, setting different weight decay values for different classes in the model is not an easy matter with PyTorch optimizers. The official documentation on <a href="https://pytorch.org/docs/stable/optim.html?highlight=optimizer#torch.optim.Optimizer" rel="noopener nofollow ugc">per-parameter options</a> means that I can set separate groups for the <code>.weight</code> attributes of the desired classes and put the rest in a separate group where weight decay is not applied.<br>
However, the difficulty is in separating the classes of the sub-modules, which are often buried inside nn.Sequential layers without names.<br>
I am currently trying to use the <code>.modules()</code> method to separate the sub-modules but I have no idea whether this is correct.</p>
<pre><code class="lang-auto">no_decay = list()
decay = list()
for m in model.modules():
  if isinstance(m, (nn.Linear, nn.Conv2d, ...)):
    decay.append(m.weight)
    no_decay.append(m.bias)
  elif hasattr(m, 'weight'):
    no_decay.append(m.weight)
  elif hasattr(m, 'bias'):
    no_decay.append(m.bias)

optimizer = optim.AdamW([{'params': no_decay, 'weight_decay', 0}, {'params': decay}, **kwargs])
</code></pre>
<p>The code above has the obvious problem that trainable parameters that are not called <code>weight</code> or <code>bias</code> are excluded completely from training.<br>
Is there any better method to do this?</p>

ptrblck: <p>Your approach seems reasonable. Iterating all named parameter of the model might be a bit easier:</p>
<pre><code class="lang-python">model = models.resnet152()

decay = dict()
no_decay = dict()
for name, m in model.named_parameters():
    print('checking {}'.format(name))
    if 'weight' in name:
        decay[name] = param
    else:
        no_decay[name] = param
    
print(decay.keys())
print(no_decay.keys())
</code></pre>

veritas: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you for the solution. However, I believe that this requires naming all layers that go inside an <code>nn.Sequential</code> layer. By default, layers within an <code>nn.Sequential</code> do not retain their names.</p>
<p>As an alternative, I have considered adding the following code.</p>
<pre><code class="lang-auto">no_decay = list()
decay = list()
for m in model.modules():
  if isinstance(m, (nn.Linear, nn.Conv2d, ...)):
    decay.append(m.weight)
    no_decay.append(m.bias)
  elif hasattr(m, 'weight'):
    no_decay.append(m.weight)
  elif hasattr(m, 'bias'):
    no_decay.append(m.bias)
for name, param in self.model.named_modules():
  if not (name.endswith('.weight') or name.endswith('.bias')):
    no_decay.append(param)
optimizer = optim.AdamW([{'params': no_decay, 'weight_decay', 0}, {'params': decay}, **kwargs])
</code></pre>
<p>The reason for extracting only the <code>weight</code> and <code>bias</code> values is that <code>.modules()</code> returns all modules, including modules that contain other modules, whereas <code>.named_parameters()</code> only returns the parameters at the very end of the recursion.</p>

ptrblck: <p><code>nn.Sequential</code> modules will add the index to the parameter names, such as <code>0.weight</code>, <code>1.weight</code> etc.</p>
<aside class="quote no-group" data-username="veritas" data-post="3" data-topic="114348">
<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/veritas/48/11018_2.png" class="avatar"> veritas:</div>
<blockquote>
<p>The reason for extracting only the <code>weight</code> and <code>bias</code> values is that <code>.modules()</code> returns all modules, including modules that contain other modules, whereas <code>.named_parameters()</code> only returns the parameters at the very end of the recursion.</p>
</blockquote>
</aside>
<p>I might misunderstand the use case, but I thought that’s exactly what you need: the actual parameters at the most inner level.</p>

veritas: <p>Hello.<br>
The reason that I used <code>.modules()</code> was to exclude normalization layers such as <code>nn.LayerNorm</code> and <code>nn.BatchNorm</code>.<br>
These layers also have parameters named <code>weight</code> and <code>bias</code> but these should be excluded from weight decay.<br>
However, within <code>nn.Sequential</code>, it is impossible to tell the type of layer from their names unless their names have been specified. This is why I use <code>.modules()</code> instead of <code>.named_parameters()</code>.</p>

veritas: <p>On second thoughts, it is much better to use Python sets.</p>
<pre><code class="lang-auto">all_params = set(model.parameters())
wd_params = set()
for m in model.modules():
  if isinstance(m, (nn.Linear, nn.Conv*)):
    wd_params.add(m.weight)
no_wd = all_params - wd_params
</code></pre>

john90: <p>Do you happen to have an alternative to using sets? This <a href="https://pytorch.org/docs/stable/optim.html?highlight=optimizer#torch.optim.Optimizer" rel="noopener nofollow ugc">PyTorch Page on optimizers</a> says explicitly that sets do not satisfy the requirements when defining the optimizer.</p>
<blockquote>
<p>Parameters need to be specified as collections that have a deterministic ordering that is consistent between runs. Examples of objects that don’t satisfy those properties are sets and iterators over values of dictionaries.</p>
</blockquote>
<p>Thanks in advance!</p>

veritas: <p><a class="mention" href="/u/john90">@john90</a>  Thank you for pointing this out! It may have been causing unknown issues in my code for a very long time.</p>
<p>I have a simple solution that does not use any undefined orderings.</p>
<pre><code class="lang-auto">import torch
from torch import nn


@torch.no_grad()
def get_wd_params(model: nn.Module):
    # Parameters must have a defined order. 
    # No sets or dictionary iterations.
    # See https://pytorch.org/docs/stable/optim.html#base-class
    # Parameters for weight decay.
    all_params = tuple(model.parameters())
    wd_params = list()
    for m in model.modules():
        if isinstance(
                m,
                (
                        nn.Linear,
                        nn.Conv1d,
                        nn.Conv2d,
                        nn.Conv3d,
                        nn.ConvTranspose1d,
                        nn.ConvTranspose2d,
                        nn.ConvTranspose3d,
                ),
        ):
            wd_params.append(m.weight)
    # Only weights of specific layers should undergo weight decay.
    no_wd_params = [p for p in all_params if p not in wd_params]
    assert len(wd_params) + len(no_wd_params) == len(all_params), "Sanity check failed."
    return wd_params, no_wd_params

</code></pre>

Jiacheng_Li: <p>Hi, I have the same problem and googled here. Below is my solution, which combines both solutions mentioned above.</p>
<pre><code class="lang-auto">@torch.no_grad()
def get_wd_params(model: nn.Module):
    decay = list()
    no_decay = list()
    for name, param in model.named_parameters():
        print('checking {}'.format(name))
        if hasattr(param,'requires_grad') and not param.requires_grad:
            continue
        if 'weight' in name and 'norm' not in name and 'bn' not in name:
            decay.append(param)
        else:
            no_decay.append(param)
    return decay, no_decay
</code></pre>

veritas: <p>Fortunately, I found that I had been converting all sets to tuples before giving them to the optimizer, so there were no problems.</p>

kaiwenli325: <p><a class="mention" href="/u/veritas">@veritas</a> Check the code here <a href="https://github.com/karpathy/minGPT/blob/3ed14b2cec0dfdad3f4b2831f2b4a86d11aef150/mingpt/model.py#L136" rel="noopener nofollow ugc">minGPT</a>. I think it is a good implementation.</p>