implement better logic for detecting weights/modules

Author: Vectorrent

pytorch_optimizer/optimizer/utils.py

@@ -205,16 +205,41 @@ def get_optimizer_parameters(
     :param wd_ban_list: List[str]. ban list not to set weight decay.
     :returns: PARAMETERS. new parameter list.
     """
+    
+    def find_fully_qualified_names(
+        model: nn.Module,
+        wd_ban_list: List[str] = ("bias", "LayerNorm.weight", "LayerNorm.bias"),
+    ):
+        names_without_wd = []
+
+        for module_name, module in model.named_modules():
+            for param_name, param in module.named_parameters(recurse=False):
+                # Full parameter name includes module and parameter names
+                full_param_name = (
+                    f"{module_name}.{param_name}" if module_name else param_name
+                )
+                # Check if any ban list substring is in the parameter name or module name
+                if (
+                    any(banned in param_name for banned in wd_ban_list)
+                    or any(banned in module_name for banned in wd_ban_list)
+                    or any(banned in module._get_name() for banned in wd_ban_list)
+                ):
+                    names_without_wd.append(full_param_name)
+
+        return names_without_wd
+
+    full_names = find_fully_qualified_names(model_or_parameter, wd_ban_list)
+
     if isinstance(model_or_parameter, nn.Module):
         model_or_parameter = list(model_or_parameter.named_parameters())
-
+    
     return [
         {
-            'params': [p for n, p in model_or_parameter if p.requires_grad and not any(nd in n for nd in wd_ban_list)],
+            'params': [p for n, p in model_or_parameter if p.requires_grad and not any(nd in n for nd in full_names)],
             'weight_decay': weight_decay,
         },
         {
-            'params': [p for n, p in model_or_parameter if p.requires_grad and any(nd in n for nd in wd_ban_list)],
+            'params': [p for n, p in model_or_parameter if p.requires_grad and any(nd in n for nd in full_names)],
             'weight_decay': 0.0,
         },
     ]

tests/test_utils.py

@@ -98,7 +98,7 @@ def test_neuron_mean_norm():
 
 def test_get_optimizer_parameters():
     model: nn.Module = Example()
-    wd_ban_list: List[str] = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+    wd_ban_list: List[str] = ['bias', 'LayerNorm.bias', 'LayerNorm.weight', 'LayerNorm']
 
     before_parameters = list(model.named_parameters())
     after_parameters = get_optimizer_parameters(model, weight_decay=1e-3, wd_ban_list=wd_ban_list)