update: add DecayParameter and NoDecayParameter

typoverflow · Sep 6, 2023 · c2d8326 · c2d8326
1 parent 30a6238
commit c2d8326
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/offlinerllib/__init__.py b/offlinerllib/__init__.py
@@ -1,2 +1,2 @@
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
diff --git a/offlinerllib/module/net/attention/base.py b/offlinerllib/module/net/attention/base.py
@@ -4,6 +4,13 @@
 from offlinerllib.module.net.attention.positional_encoding import BasePosEncoding
 
 
+class DecayParameter(nn.Parameter):
+    pass
+
+class NoDecayParameter(nn.Parameter):
+    pass
+
+
 class BaseTransformer(nn.Module):
     def __init__(self, *args, **kwargs):
         super().__init__()
@@ -31,6 +38,10 @@ def configure_params(self):
                     no_decay.add(fpn)
                 elif pn.endswith("weight") and isinstance(m, whitelist_weight_modules):
                     decay.add(fpn)
+                elif isinstance(p, DecayParameter):
+                    decay.add(fpn)
+                elif isinstance(p, NoDecayParameter):
+                    no_decay.add(fpn)
 
         # validate that we considered every parameter
         param_dict = {pn: p for pn, p in self.named_parameters()}