In many of the papers and blogs that I read, for example, the recent NFNet 20 paper, the authors emphasize the importance of only including the convolution & 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.
However, setting different weight decay values for different classes in the model is not an easy matter with PyTorch optimizers. The official documentation on per-parameter options 9 means that I can set separate groups for the .weight attributes of the desired classes and put the rest in a separate group where weight decay is not applied.
However, the difficulty is in separating the classes of the sub-modules, which are often buried inside nn.Sequential layers without names.
I am currently trying to use the .modules() method to separate the sub-modules but I have no idea whether this is correct.
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])

The code above has the obvious problem that trainable parameters that are not called weight or bias are excluded completely from training.
Is there any better method to do this?