feat(optim): WeightDecayChain

Author: ClashLuke

setup.py

@@ -10,7 +10,7 @@
     name='truegrad',
     license='BSD',
     description='PyTorch interface for TrueGrad-AdamW',
-    version='3.1.1',
+    version='4.0.0',
     long_description=README,
     url='https://github.com/clashluke/truegrad',
     packages=setuptools.find_packages(),

truegrad/optim.py

@@ -1,3 +1,4 @@
+import functools
 import warnings
 from typing import Tuple, Union, List, Dict, Any, Optional
 
@@ -6,6 +7,94 @@
 from torch.nn import Parameter
 
 
+class WeightDecayBase:
+    def __init__(self):
+        pass
+
+    def __call__(self, mod: torch.optim.Optimizer, p: torch.Tensor, idx: int):
+        return p
+
+
+class WeightDecayChain:
+    def __init__(self, *operands: WeightDecayBase):
+        self.operands = operands
+
+    def __call__(self, mod: torch.optim.Optimizer):
+        idx = 0
+        for group in mod.param_groups:
+            for p in group["params"]:
+                p.data.add(functools.reduce(lambda x, f: f(mod, x, idx), self.operands, p),
+                           alpha=-group["lr"] * group["weight_decay"])
+                idx += 1
+
+
+class LpWeightDecay(WeightDecayBase):
+    def __init__(self, power: float):
+        self.power = power
+
+    def __call__(self, mod: torch.optim.Optimizer, p: Tensor, idx: int):
+        return p.abs().pow(self.power) * p.sign()
+
+
+class L1WeightDecay(LpWeightDecay):
+    def __init__(self):
+        super().__init__(0)
+
+
+class L2WeightDecay(LpWeightDecay):
+    def __init__(self):
+        super().__init__(1)
+
+
+def _param_iterator(mod: torch.optim.Optimizer):
+    yield from (p.detach().clone() for group in mod.param_groups for p in group["params"])
+
+
+class WeightDecayToValue(WeightDecayBase):
+    def __init__(self):
+        self.target_values: List[Tensor] = ...
+        self.global_step = 0
+
+    def _on_step_start(self, mod: torch.optim.Optimizer):
+        pass
+
+    def _on_global_start(self, mod: torch.optim.Optimizer):
+        pass
+
+    def _preprocess(self, target: Tensor):
+        return target
+
+    def __call__(self, mod: torch.optim.Optimizer, p: Tensor, idx: int):
+        if idx == 0:
+            if self.global_step == 0:
+                self._on_global_start(mod)
+            self._on_step_start(mod)
+            self.global_step += 1
+        return p - self._preprocess(self.target_values[idx])
+
+
+class WeightDecayToInit(WeightDecayToValue):
+    def _on_global_start(self, mod: torch.optim.Optimizer):
+        self.target_values = list(_param_iterator(mod))
+
+
+class WeightDecayToEMA(WeightDecayToInit):
+    def __init__(self, beta: float = 0.999):
+        super().__init__()
+        self.beta = beta
+
+    def _on_global_start(self, mod: torch.optim.Optimizer):
+        self.target_values = [torch.zeros_like(x) for x in _param_iterator(mod)]
+
+    def _on_step_start(self, mod: torch.optim.Optimizer):
+        self.global_step += 1
+        for v, p in zip(self.target_values, _param_iterator(mod)):
+            v.mul_(self.beta).add_(p, alpha=1 - self.beta)
+
+    def _preprocess(self, target: Tensor):
+        return target / (1 - self.beta ** self.global_step)
+
+
 def ema_(base: Tensor, update: Tensor, beta: float, step: Optional[int] = None):
     base.mul_(beta).add_(update, alpha=1 - beta)
     if step is None:
@@ -31,12 +120,18 @@ def decay_weight_(state: Dict[str, Any], param: torch.nn.Parameter, group: Dict[
         param.mul_(1 - group["weight_decay"] * group["lr"])
 
 
+def _default_decay(weight_decay_cls: Optional[WeightDecayChain]) -> WeightDecayChain:
+    if weight_decay_cls is None:
+        return WeightDecayChain(L2WeightDecay())
+    return weight_decay_cls
+
+
 class OptimizerOptimizer(torch.optim.Optimizer):
     def __init__(self, params, inner_optimizer: torch.optim.Optimizer, learning_rate_learning_rate: float = 1,
-                 weight_decay: float = 0, decay_to_init: bool = False):
-        self.learning_rate_learning_rate = learning_rate_learning_rate
-
+                 weight_decay: float = 0, weight_decay_cls: Optional[WeightDecayChain] = None):
         self.inner_optimizer = inner_optimizer
+        self.learning_rate_learning_rate = learning_rate_learning_rate
+        self.weight_decay_cls = _default_decay(weight_decay_cls)
         param_groups = self.inner_optimizer.param_groups
         self.inner_optimizer.param_groups = []
         for group in param_groups:
@@ -45,14 +140,16 @@ def __init__(self, params, inner_optimizer: torch.optim.Optimizer, learning_rate
                 group["params"] = [param]
                 self.inner_optimizer.param_groups.append(group)
 
-        super(OptimizerOptimizer, self).__init__(params, {"weight_decay": weight_decay, "decay_to_init": decay_to_init})
+        super(OptimizerOptimizer, self).__init__(params, {"weight_decay": weight_decay})
 
     @torch.no_grad()
     def step(self, closure=None):
         loss = None
         if closure is not None:
             loss = closure()
 
+        self.weight_decay_cls(self)
+
         for group in self.param_groups:
             for p in group['params']:
                 state = self.state[p]
@@ -80,10 +177,11 @@ def step(self, closure=None):
 
 
 class Sign(torch.optim.Optimizer):
-    def __init__(self, params, base: torch.optim.Optimizer, lr: float = 1, weight_decay: float = 0,
-                 decay_to_init: bool = False, eps: float = 1e-12, graft_to_self: bool = True):
-        super().__init__(params, {"weight_decay": weight_decay, "decay_to_init": decay_to_init, "lr": lr, "eps": eps,
-                                  "graft_to_self": graft_to_self})
+    def __init__(self, params, base: torch.optim.Optimizer, lr: float = 1, weight_decay: float = 0, eps: float = 1e-12,
+                 graft_to_self: bool = True, weight_decay_cls: Optional[WeightDecayChain] = None):
+        self.weight_decay_cls = _default_decay(weight_decay_cls)
+
+        super().__init__(params, {"weight_decay": weight_decay, "lr": lr, "eps": eps, "graft_to_self": graft_to_self})
         self.base = base
 
     @torch.no_grad()
@@ -94,14 +192,8 @@ def step(self, closure=None):
             with torch.enable_grad():
                 loss = closure()
 
-        params_flat = []
-        for group in self.param_groups:
-            for p in group["params"]:
-                params_flat.append(p)
-                decay_weight_(self.state[p], p, group)
-
-        params_flat = [torch.clone(p.detach()) for p in params_flat]
-
+        self.weight_decay_cls(self)
+        params_flat = list(_param_iterator(self))
         self.base.step()
 
         for group in self.param_groups:
@@ -150,10 +242,12 @@ class Graft(torch.optim.Optimizer):
     """
 
     def __init__(self, params, magnitude: torch.optim.Optimizer, direction: torch.optim.Optimizer,
-                 weight_decay: float = 0, decay_to_init: bool = False, eps: float = 1e-12, lr: float = 1):
-        super().__init__(params, {"weight_decay": weight_decay, "decay_to_init": decay_to_init, "lr": lr, "eps": eps})
+                 weight_decay: float = 0, eps: float = 1e-12, lr: float = 1,
+                 weight_decay_cls: Optional[WeightDecayChain] = None):
+        super().__init__(params, {"weight_decay": weight_decay, "lr": lr, "eps": eps})
         self.magnitude = magnitude
         self.direction = direction
+        self.weight_decay_cls = _default_decay(weight_decay_cls)
 
     @torch.no_grad()
     def step(self, closure=None):
@@ -163,13 +257,8 @@ def step(self, closure=None):
             with torch.enable_grad():
                 loss = closure()
 
-        params_flat = []
-        for group in self.param_groups:
-            for p in group["params"]:
-                params_flat.append(p)
-                decay_weight_(self.state[p], p, group)
-
-        original_params = [torch.clone(p.detach()) for p in params_flat]
+        self.weight_decay_cls(self)
+        original_params = list(_param_iterator(self))
 
         self.magnitude.step()
         magnitudes_flat = []
@@ -194,21 +283,16 @@ class TrueGrad(torch.optim.Optimizer):
     base_statistics: List[str] = []
     shared_statistics: List[str] = []
 
-    def __init__(self, params, lr: float = 1e-3,
-                 betas: List[float] = (),
-                 eps: float = 1e-12,
-                 weight_decay: float = 1e-2,
-                 graft: bool = True,
-                 decay_to_init: bool = False,
-                 default_to_baseline: bool = False,
-                 enforce_baseline: bool = False):
+    def __init__(self, params, lr: float = 1e-3, betas: List[float] = (), eps: float = 1e-12,
+                 weight_decay: float = 1e-2, graft: bool = True, default_to_baseline: bool = False,
+                 enforce_baseline: bool = False, weight_decay_cls: Optional[WeightDecayChain] = None):
         defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, graft=graft,
-                        decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
-                        enforce_baseline=enforce_baseline)
+                        default_to_baseline=default_to_baseline, enforce_baseline=enforce_baseline)
         super(TrueGrad, self).__init__(params, defaults)
+        self.weight_decay_cls = _default_decay(weight_decay_cls)
 
-    def _inner(self, step: int, p: Parameter, group: Dict[str, Any], **kwargs: Tensor
-               ) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
+    def _inner(self, step: int, p: Parameter, group: Dict[str, Any], **kwargs: Tensor) -> Tuple[
+        Optional[Tensor], Optional[Tensor], float]:
         raise NotImplementedError
 
     @torch.no_grad()
@@ -245,12 +329,7 @@ def step(self, closure=None):
                 step_t = state['step']
                 step_t += 1
 
-                # Perform stepweight decay
-                decay = group['lr'] * group['weight_decay']
-                if group["decay_to_init"]:
-                    p.add_(state["init"] - p, alpha=decay)
-                else:
-                    p.mul_(1 - decay)
+                self.weight_decay_cls(self)
 
                 step = step_t.item()
 
@@ -275,28 +354,18 @@ class TGAdamW(TrueGrad):
 
     def __init__(self, params, lr: float = 1e-3,
                  betas: Union[Tuple[float, float], Tuple[float, float, float]] = (0.9, 0.999, 0.999),
-                 eps: float = 1e-12,
-                 weight_decay: float = 1e-2,
-                 graft: bool = True,
-                 decay_to_init: bool = False,
-                 default_to_adam: bool = None,
-                 default_to_baseline: bool = None,
-                 enforce_baseline: bool = False):
-        if default_to_baseline is None:
-            default_to_baseline = default_to_adam
-        elif default_to_adam is not None:
-            raise ValueError("Can't set both default_to_baseline and default_to_adam, as both map to the same argument")
-        if default_to_adam is not None:
-            warnings.warn("default_to_adam is deprecated and will be replaced by default_to_baseline in April 2023")
+                 eps: float = 1e-12, weight_decay: float = 1e-2, graft: bool = True,
+                 default_to_baseline: bool = None, enforce_baseline: bool = False,
+                 weight_decay_cls: Optional[WeightDecayChain] = None):
         if default_to_baseline is None:
             default_to_baseline = False
         super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, graft=graft,
-                         decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
-                         enforce_baseline=enforce_baseline)
+                         default_to_baseline=default_to_baseline, enforce_baseline=enforce_baseline,
+                         weight_decay_cls=weight_decay_cls)
 
     def _inner(self, step: int, p: Parameter, group: Dict[str, Any], exp_avg: Tensor,
-               exp_avg_sq: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None
-               ) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
+               exp_avg_sq: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None) -> Tuple[
+        Optional[Tensor], Optional[Tensor], float]:
         if len(group["betas"]) == 2:
             (beta1, beta2), (_, beta3) = group["betas"], group["betas"]
         else:
@@ -317,21 +386,17 @@ class TGLaProp(TrueGrad):
     base_statistics: List[str] = ["exp_avg", "exp_avg_sq"]
 
     def __init__(self, params, lr: float = 1e-3,
-                 betas: Union[Tuple[float, float], Tuple[float, float, float, float]] = (0.9, 0.99),
-                 eps: float = 1e-12,
-                 weight_decay: float = 1e-2,
-                 graft: bool = True,
-                 decay_to_init: bool = False,
-                 default_to_baseline: bool = False,
-                 enforce_baseline: bool = False):
+                 betas: Union[Tuple[float, float], Tuple[float, float, float, float]] = (0.9, 0.99), eps: float = 1e-12,
+                 weight_decay: float = 1e-2, graft: bool = True, decay_to_init: bool = False,
+                 default_to_baseline: bool = False, enforce_baseline: bool = False,
+                 weight_decay_cls: Optional[WeightDecayChain] = None):
         super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, graft=graft,
-                         decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
-                         enforce_baseline=enforce_baseline)
+                         default_to_baseline=default_to_baseline, enforce_baseline=enforce_baseline,
+                         weight_decay_cls=weight_decay_cls)
 
-    def _inner(self, step: int, p: Parameter, group: Dict[str, Any],
-               exp_avg: Optional[Tensor] = None, exp_avg_sq: Optional[Tensor] = None,
-               exp_avg_true: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None
-               ) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
+    def _inner(self, step: int, p: Parameter, group: Dict[str, Any], exp_avg: Optional[Tensor] = None,
+               exp_avg_sq: Optional[Tensor] = None, exp_avg_true: Optional[Tensor] = None,
+               exp_avg_true_sq: Optional[Tensor] = None) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
         if len(group["betas"]) == 2:
             (beta1, beta2), (beta3, beta4) = group["betas"], group["betas"]
         else:
@@ -362,21 +427,16 @@ class TGRMSProp(TrueGrad):
     true_statistics: List[str] = ["exp_avg_true_sq"]
     base_statistics: List[str] = ["exp_avg_sq"]
 
-    def __init__(self, params, lr: float = 1e-3,
-                 betas: Union[float, Tuple[float], Tuple[float, float]] = (0.9,),
-                 eps: float = 1e-12,
-                 weight_decay: float = 1e-2,
-                 graft: bool = True,
-                 decay_to_init: bool = False,
-                 default_to_baseline: bool = False,
-                 enforce_baseline: bool = False):
+    def __init__(self, params, lr: float = 1e-3, betas: Union[float, Tuple[float], Tuple[float, float]] = (0.9,),
+                 eps: float = 1e-12, weight_decay: float = 1e-2, graft: bool = True,
+                 default_to_baseline: bool = False, enforce_baseline: bool = False,
+                 weight_decay_cls: Optional[WeightDecayChain] = None):
         super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, graft=graft,
-                         decay_to_init=decay_to_init, default_to_baseline=default_to_baseline,
-                         enforce_baseline=enforce_baseline)
+                         default_to_baseline=default_to_baseline, enforce_baseline=enforce_baseline,
+                         weight_decay_cls=weight_decay_cls)
 
-    def _inner(self, step: int, p: Parameter, group: Dict[str, Any],
-               exp_avg_sq: Optional[Tensor] = None, exp_avg_true_sq: Optional[Tensor] = None
-               ) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
+    def _inner(self, step: int, p: Parameter, group: Dict[str, Any], exp_avg_sq: Optional[Tensor] = None,
+               exp_avg_true_sq: Optional[Tensor] = None) -> Tuple[Optional[Tensor], Optional[Tensor], float]:
         if isinstance(group["betas"], float):
             beta1 = beta2 = group["betas"]
         elif len(group["betas"]) == 1: