Merge pull request #35 from kozistr/feature/wd-utils

[Feature] weight decay ban list

Author: kozistr

lint.py

@@ -14,7 +14,7 @@ def get_configuration() -> Namespace:
     parser.add_argument(
         '-t',
         '--threshold',
-        default=9.75,
+        default=9.9,
         type=float,
     )
 

pytorch_optimizer/__init__.py

@@ -16,5 +16,6 @@
 from pytorch_optimizer.ranger21 import Ranger21
 from pytorch_optimizer.sam import SAM
 from pytorch_optimizer.sgdp import SGDP
+from pytorch_optimizer.utils import get_optimizer_parameters, normalize_gradient, unit_norm
 
-__VERSION__ = '0.1.0'
+__VERSION__ = '0.1.1'

pytorch_optimizer/utils.py

@@ -1,6 +1,7 @@
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
+from torch import nn
 
 from pytorch_optimizer.types import PARAMETERS
 
@@ -41,3 +42,17 @@ def unit_norm(x: torch.Tensor, norm: float = 2.0) -> torch.Tensor:
         dim = tuple(range(1, x_len))
 
     return x.norm(dim=dim, keepdim=keep_dim, p=norm)
+
+
+def get_optimizer_parameters(
+    model: nn.Module, weight_decay: float, wd_ban_list: List[str] = ('bias', 'LayerNorm.bias', 'LayerNorm.weight')
+) -> PARAMETERS:
+    param_optimizer: List[Tuple[str, nn.Parameter]] = list(model.named_parameters())
+
+    return [
+        {
+            'params': [p for n, p in param_optimizer if not any(nd in n for nd in wd_ban_list)],
+            'weight_decay': weight_decay,
+        },
+        {'params': [p for n, p in param_optimizer if any(nd in n for nd in wd_ban_list)], 'weight_decay': 0.0},
+    ]