update: test_get_optimizer_parameters

Author: kozistr

tests/test_utils.py

@@ -1,7 +1,26 @@
+from typing import List
+
 import numpy as np
 import torch
+from torch import nn
+
+from pytorch_optimizer.utils import (
+    clip_grad_norm,
+    get_optimizer_parameters,
+    has_overflow,
+    normalize_gradient,
+    unit_norm,
+)
+
+
+class Example(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = nn.Linear(1, 1)
+        self.norm1 = nn.LayerNorm(1)
 
-from pytorch_optimizer.utils import clip_grad_norm, has_overflow, normalize_gradient, unit_norm
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.norm1(self.fc1(x))
 
 
 def test_has_overflow():
@@ -43,3 +62,16 @@ def test_unit_norm():
     np.testing.assert_approx_equal(unit_norm(x.view(1, 10)).numpy(), 16.8819, significant=4)
     np.testing.assert_approx_equal(unit_norm(x.view(1, 10, 1, 1)).numpy(), 16.8819, significant=4)
     np.testing.assert_approx_equal(unit_norm(x.view(1, 10, 1, 1, 1, 1)).numpy(), 16.8819, significant=4)
+
+
+def test_get_optimizer_parameters():
+    model: nn.Module = Example()
+    wd_ban_list: List[str] = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
+
+    before_parameters = list(model.named_parameters())
+    after_parameters = get_optimizer_parameters(model, weight_decay=1e-3, wd_ban_list=wd_ban_list)
+
+    for before, after in zip(before_parameters, after_parameters):
+        layer_name: str = before[0]
+        if layer_name.find('bias') != -1 or layer_name in wd_ban_list:
+            assert after['weight_decay'] == 0.0