torchzero 0.3.10__py3-none-any.whl → 0.3.13__py3-none-any.whl

torchzero/modules/variance_reduction/svrg.py

@@ -0,0 +1,208 @@
+import warnings
+from functools import partial
+
+import torch
+
+from ...core.module import Module
+from ...utils import tofloat
+
+
+def _reset_except_self(optimizer, var, self: Module):
+    for m in optimizer.unrolled_modules:
+        if m is not self:
+            m.reset()
+
+class SVRG(Module):
+    """Stochastic variance reduced gradient method (SVRG).
+
+    To use, put SVRG as the first module, it can be used with any other modules.
+    To reduce variance of a gradient estimator, put the gradient estimator before SVRG.
+
+    First it uses first ``accum_steps`` batches to compute full gradient at initial
+    parameters using gradient accumulation, the model will not be updated during this.
+
+    Then it performs ``svrg_steps`` SVRG steps, each requires two forward and backward passes.
+
+    After ``svrg_steps``, it goes back to full gradient computation step step.
+
+    As an alternative to gradient accumulation you can pass "full_closure" argument to the ``step`` method,
+    which should compute full gradients, set them to ``.grad`` attributes of the parameters,
+    and return full loss.
+
+    Args:
+        svrg_steps (int): number of steps before calculating full gradient. This can be set to length of the dataloader.
+        accum_steps (int | None, optional):
+            number of steps to accumulate the gradient for. Not used if "full_closure" is passed to the ``step`` method. If None, uses value of ``svrg_steps``. Defaults to None.
+        reset_before_accum (bool, optional):
+            whether to reset all other modules when re-calculating full gradient. Defaults to True.
+        svrg_loss (bool, optional):
+            whether to replace loss with SVRG loss (calculated by same formula as SVRG gradient). Defaults to True.
+        alpha (float, optional):
+            multiplier to ``g_full(x_0) - g_batch(x_0)`` term, can be annealed linearly from 1 to 0 as suggested in https://arxiv.org/pdf/2311.05589#page=6
+
+    ## Examples:
+    SVRG-LBFGS
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SVRG(len(dataloader)),
+        tz.m.LBFGS(),
+        tz.m.Backtracking(),
+    )
+    ```
+
+    For extra variance reduction one can use Online versions of algorithms, although it won't always help.
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SVRG(len(dataloader)),
+        tz.m.Online(tz.m.LBFGS()),
+        tz.m.Backtracking(),
+    )
+
+    Variance reduction can also be applied to gradient estimators.
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.SPSA(),
+        tz.m.SVRG(100),
+        tz.m.LR(1e-2),
+    )
+    ```
+    ## Notes
+
+    The SVRG gradient is computed as ``g_b(x) - alpha * g_b(x_0) - g_f(x0.)``, where:
+    - ``x`` is current parameters
+    - ``x_0`` is initial parameters, where full gradient was computed
+    - ``g_b`` refers to mini-batch gradient at ``x`` or ``x_0``
+    - ``g_f`` refers to full gradient at ``x_0``.
+
+    The SVRG loss is computed using the same formula.
+    """
+    def __init__(self, svrg_steps: int, accum_steps: int | None = None, reset_before_accum:bool=True, svrg_loss:bool=True, alpha:float=1):
+        defaults = dict(svrg_steps = svrg_steps, accum_steps=accum_steps, reset_before_accum=reset_before_accum, svrg_loss=svrg_loss, alpha=alpha)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        params = var.params
+        closure = var.closure
+        assert closure is not None
+
+        if "full_grad" not in self.global_state:
+
+            # -------------------------- calculate full gradient ------------------------- #
+            if "full_closure" in var.storage:
+                full_closure = var.storage['full_closure']
+                with torch.enable_grad():
+                    full_loss = full_closure()
+                    if all(p.grad is None for p in params):
+                        warnings.warn("all gradients are None after evaluating full_closure.")
+
+                    full_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    self.global_state["full_loss"] = full_loss
+                    self.global_state["full_grad"] = full_grad
+                    self.global_state['x_0'] = [p.clone() for p in params]
+
+                # current batch will be used for svrg update
+
+            else:
+                # accumulate gradients over n steps
+                accum_steps = self.defaults['accum_steps']
+                if accum_steps is None: accum_steps = self.defaults['svrg_steps']
+
+                current_accum_step = self.global_state.get('current_accum_step', 0) + 1
+                self.global_state['current_accum_step'] = current_accum_step
+
+                # accumulate grads
+                accumulator = self.get_state(params, 'accumulator')
+                grad = var.get_grad()
+                torch._foreach_add_(accumulator, grad)
+
+                # accumulate loss
+                loss_accumulator = self.global_state.get('loss_accumulator', 0)
+                loss_accumulator += tofloat(var.loss)
+                self.global_state['loss_accumulator'] = loss_accumulator
+
+                # on nth step, use the accumulated gradient
+                if current_accum_step >= accum_steps:
+                    torch._foreach_div_(accumulator, accum_steps)
+                    self.global_state["full_grad"] = accumulator
+                    self.global_state["full_loss"] = loss_accumulator / accum_steps
+
+                    self.global_state['x_0'] = [p.clone() for p in params]
+                    self.clear_state_keys('accumulator')
+                    del self.global_state['current_accum_step']
+
+                # otherwise skip update until enough grads are accumulated
+                else:
+                    var.update = None
+                    var.stop = True
+                    var.skip_update = True
+                    return var
+
+
+        svrg_steps = self.defaults['svrg_steps']
+        current_svrg_step = self.global_state.get('current_svrg_step', 0) + 1
+        self.global_state['current_svrg_step'] = current_svrg_step
+
+        # --------------------------- SVRG gradient closure -------------------------- #
+        x0 = self.global_state['x_0']
+        gf_x0 = self.global_state["full_grad"]
+        ff_x0 = self.global_state['full_loss']
+        use_svrg_loss = self.defaults['svrg_loss']
+        alpha = self.get_settings(params, 'alpha')
+        alpha_0 = alpha[0]
+        if all(a == 1 for a in alpha): alpha = None
+
+        def svrg_closure(backward=True):
+            # g_b(x) - α * (g_f(x_0) - g_b(x_0)) and same for loss
+            with torch.no_grad():
+                x = [p.clone() for p in params]
+
+                if backward:
+                    # f and g at x
+                    with torch.enable_grad(): fb_x = closure()
+                    gb_x = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+
+                    # f and g at x_0
+                    torch._foreach_copy_(params, x0)
+                    with torch.enable_grad(): fb_x0 = closure()
+                    gb_x0 = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+                    torch._foreach_copy_(params, x)
+
+                    # g_svrg = gb_x - alpha * (gf_x0 - gb_x0)
+                    correction = torch._foreach_sub(gb_x0, gf_x0)
+                    if alpha is not None: torch._foreach_mul_(correction, alpha)
+                    g_svrg = torch._foreach_sub(gb_x, correction)
+
+                    f_svrg = fb_x - alpha_0 * (fb_x0 - ff_x0)
+                    for p, g in zip(params, g_svrg):
+                        p.grad = g
+
+                    if use_svrg_loss: return f_svrg
+                    return fb_x
+
+            # no backward
+            if use_svrg_loss:
+                fb_x = closure(False)
+                torch._foreach_copy_(params, x0)
+                fb_x0 = closure(False)
+                torch._foreach_copy_(params, x)
+                f_svrg = fb_x - alpha_0 * (fb_x0 - ff_x0)
+                return f_svrg
+
+            return closure(False)
+
+        var.closure = svrg_closure
+
+        # --- after svrg_steps steps reset so that new full gradient is calculated on next step --- #
+        if current_svrg_step >= svrg_steps:
+            del self.global_state['current_svrg_step']
+            del self.global_state['full_grad']
+            del self.global_state['full_loss']
+            del self.global_state['x_0']
+            if self.defaults['reset_before_accum']:
+                var.post_step_hooks.append(partial(_reset_except_self, self=self))
+
+        return var

torchzero/modules/weight_decay/__init__.py

@@ -1 +1 @@
-from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_, NormalizedWeightDecay
+from .weight_decay import WeightDecay, DirectWeightDecay, decay_weights_, RelativeWeightDecay

torchzero/modules/weight_decay/weight_decay.py

@@ -4,7 +4,7 @@ from typing import Literal
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states, Metrics
 
 
 @torch.no_grad
@@ -14,7 +14,7 @@ def weight_decay_(
     weight_decay: float | NumberList,
     ord: int = 2
 ):
-    """returns `grad_`."""
+    """modifies in-place and returns ``grad_``."""
     if ord == 1: return grad_.add_(params.sign().mul_(weight_decay))
     if ord == 2: return grad_.add_(params.mul(weight_decay))
     if ord - 1 % 2 != 0: return grad_.add_(params.pow(ord-1).mul_(weight_decay))
@@ -22,34 +22,113 @@ def weight_decay_(
 
 
 class WeightDecay(Transform):
+    """Weight decay.
+
+    Args:
+        weight_decay (float): weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+        target (Target, optional): what to set on var. Defaults to 'update'.
+
+    ### Examples:
+
+    Adam with non-decoupled weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.WeightDecay(1e-3),
+        tz.m.Adam(),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with decoupled weight decay that still scales with learning rate
+    ```python
+
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.WeightDecay(1e-3),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with fully decoupled weight decay that doesn't scale with learning rate
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.LR(1e-3),
+        tz.m.WeightDecay(1e-6)
+    )
+    ```
+
+    """
     def __init__(self, weight_decay: float, ord: int = 2, target: Target = 'update'):
+
         defaults = dict(weight_decay=weight_decay, ord=ord)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         weight_decay = NumberList(s['weight_decay'] for s in settings)
         ord = settings[0]['ord']
 
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay, ord)
 
-class NormalizedWeightDecay(Transform):
+class RelativeWeightDecay(Transform):
+    """Weight decay relative to the mean absolute value of update, gradient or parameters depending on value of ``norm_input`` argument.
+
+    Args:
+        weight_decay (float): relative weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+        norm_input (str, optional):
+            determines what should weight decay be relative to. "update", "grad" or "params".
+            Defaults to "update".
+        metric (Ords, optional):
+            metric (norm, etc) that weight decay should be relative to.
+            defaults to 'mad' (mean absolute deviation).
+        target (Target, optional): what to set on var. Defaults to 'update'.
+
+    ### Examples:
+
+    Adam with non-decoupled relative weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.RelativeWeightDecay(1e-1),
+        tz.m.Adam(),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with decoupled relative weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.RelativeWeightDecay(1e-1),
+        tz.m.LR(1e-3)
+    )
+    ```
+    """
     def __init__(
         self,
         weight_decay: float = 0.1,
-        ord: int = 2,
+        ord: int  = 2,
         norm_input: Literal["update", "grad", "params"] = "update",
+        metric: Metrics = 'mad',
         target: Target = "update",
     ):
-        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input)
+        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input, metric=metric)
         super().__init__(defaults, uses_grad=norm_input == 'grad', target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         weight_decay = NumberList(s['weight_decay'] for s in settings)
 
         ord = settings[0]['ord']
         norm_input = settings[0]['norm_input']
+        metric = settings[0]['metric']
 
         if norm_input == 'update': src = TensorList(tensors)
         elif norm_input == 'grad':
@@ -60,8 +139,7 @@ class NormalizedWeightDecay(Transform):
         else:
             raise ValueError(norm_input)
 
-        norm = src.global_vector_norm(ord)
-
+        norm = src.global_metric(metric)
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay * norm, ord)
 
 
@@ -72,7 +150,12 @@ def decay_weights_(params: Iterable[torch.Tensor], weight_decay: float | NumberL
     weight_decay_(params, params, -weight_decay, ord)
 
 class DirectWeightDecay(Module):
-    """directly decays weights in-place"""
+    """Directly applies weight decay to parameters.
+
+    Args:
+        weight_decay (float): weight decay scale.
+        ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
+    """
     def __init__(self, weight_decay: float, ord: int = 2,):
         defaults = dict(weight_decay=weight_decay, ord=ord)
         super().__init__(defaults)
@@ -80,7 +163,7 @@ class DirectWeightDecay(Module):
     @torch.no_grad
     def step(self, var):
         weight_decay = self.get_settings(var.params, 'weight_decay', cls=NumberList)
-        ord = self.settings[var.params[0]]['ord']
+        ord = self.defaults['ord']
 
         decay_weights_(var.params, weight_decay, ord)
         return var

torchzero/modules/wrappers/optim_wrapper.py

@@ -7,7 +7,35 @@ from ...utils import Params, _copy_param_groups, _make_param_groups
 
 
 class Wrap(Module):
-    """Custom param groups are supported only by `set_param_groups`. Settings passed to Modular will be ignored."""
+    """
+    Wraps a pytorch optimizer to use it as a module.
+
+    .. note::
+        Custom param groups are supported only by `set_param_groups`, settings passed to Modular will be ignored.
+
+    Args:
+        opt_fn (Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer):
+            function that takes in parameters and returns the optimizer, for example :code:`torch.optim.Adam`
+            or :code:`lambda parameters: torch.optim.Adam(parameters, lr=1e-3)`
+        *args:
+        **kwargs:
+            Extra args to be passed to opt_fn. The function is called as :code:`opt_fn(parameters, *args, **kwargs)`.
+
+    Example:
+        wrapping pytorch_optimizer.StableAdamW
+
+        .. code-block:: py
+
+            from pytorch_optimizer import StableAdamW
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Wrap(StableAdamW, lr=1),
+                tz.m.Cautious(),
+                tz.m.LR(1e-2)
+            )
+
+
+    """
     def __init__(self, opt_fn: Callable[..., torch.optim.Optimizer] | torch.optim.Optimizer, *args, **kwargs):
         super().__init__()
         self._opt_fn = opt_fn

torchzero/modules/zeroth_order/__init__.py

@@ -0,0 +1 @@
+from .cd import CD, CCD, CCDLS