torchzero 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

torchzero/modules/momentum/matrix_momentum.py

@@ -2,123 +2,192 @@ from typing import Literal
 
 import torch
 
-from ...core import Module, apply
+from ...core import Module, apply_transform, Chainable
 from ...utils import NumberList, TensorList, as_tensorlist
 from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
 
 class MatrixMomentum(Module):
+    """Second order momentum method.
+
+    Matrix momentum is useful for convex objectives, also for some reason it has very really good generalization on elastic net logistic regression.
+
+    .. note::
+        :code:`mu` needs to be tuned very carefully. It is supposed to be smaller than (1/largest eigenvalue), otherwise this will be very unstable.
+
+    .. note::
+        I have devised an adaptive version of this - :code:`tz.m.AdaptiveMatrixMomentum`, and it works well
+        without having to tune :code:`mu`.
+
+    .. note::
+        In most cases MatrixMomentum should be the first module in the chain because it relies on autograd.
+
+    .. note::
+        This module requires the a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+    Args:
+        mu (float, optional): this has a similar role to (1 - beta) in normal momentum. Defaults to 0.1.
+        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
+        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
+
+    Reference:
+        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
     """
-    May be useful for ill conditioned stochastic quadratic objectives but I need to test this.
-    Evaluates hessian vector product on each step (via finite difference or autograd).
 
-    `mu` is supposed to be smaller than (1/largest eigenvalue), otherwise this will be very unstable.
-
-    Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
-    """
-    def __init__(self, mu=0.1, beta:float=1, hvp_mode: Literal['autograd', 'forward', 'central'] = 'forward', h=1e-3, hvp_tfm=None):
-        defaults = dict(mu=mu, beta=beta, hvp_mode=hvp_mode, h=h)
+    def __init__(
+        self,
+        mu=0.1,
+        beta: float = 1,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h: float = 1e-3,
+        hvp_tfm: Chainable | None = None,
+    ):
+        defaults = dict(mu=mu, beta=beta, hvp_method=hvp_method, h=h)
         super().__init__(defaults)
 
         if hvp_tfm is not None:
             self.set_child('hvp_tfm', hvp_tfm)
 
-    @torch.no_grad
-    def step(self, vars):
-        assert vars.closure is not None
-        prev_update = self.get_state('prev_update', params=vars.params, cls=TensorList)
-        hvp_mode = self.settings[vars.params[0]]['hvp_mode']
-        h = self.settings[vars.params[0]]['h']
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_update')
 
-        mu,beta = self.get_settings('mu','beta', params=vars.params, cls=NumberList)
-
-        if hvp_mode == 'autograd':
-            with torch.enable_grad():
-                grad = vars.get_grad(create_graph=True)
-                hvp_ = TensorList(hvp(vars.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
+    @torch.no_grad
+    def update(self, var):
+        assert var.closure is not None
+        prev_update = self.get_state(var.params, 'prev_update')
+        hvp_method = self.settings[var.params[0]]['hvp_method']
+        h = self.settings[var.params[0]]['h']
 
-        elif hvp_mode == 'forward':
-            vars.get_grad()
-            l, hvp_ = hvp_fd_forward(vars.closure, vars.params, vec=prev_update, g_0=vars.grad, h=h, normalize=True)
-            if vars.loss_approx is None: vars.loss_approx = l
+        Hvp, _ = self.Hvp(prev_update, at_x0=True, var=var, rgrad=None, hvp_method=hvp_method, h=h, normalize=True, retain_grad=False)
+        Hvp = [t.detach() for t in Hvp]
 
-        elif hvp_mode == 'central':
-            l, hvp_ = hvp_fd_central(vars.closure, vars.params, vec=prev_update, h=h, normalize=True)
-            if vars.loss_approx is None: vars.loss_approx = l
+        if 'hvp_tfm' in self.children:
+            Hvp = TensorList(apply_transform(self.children['hvp_tfm'], Hvp, params=var.params, grads=var.grad, var=var))
 
-        else:
-            raise ValueError(hvp_mode)
+        self.store(var.params, "Hvp", Hvp)
 
-        if 'hvp_tfm' in self.children:
-            hvp_ = TensorList(apply(self.children['hvp_tfm'], hvp_, params=vars.params, grads=vars.grad, vars=vars))
 
-        update = TensorList(vars.get_update())
+    @torch.no_grad
+    def apply(self, var):
+        update = TensorList(var.get_update())
+        Hvp, prev_update = self.get_state(var.params, 'Hvp', 'prev_update', cls=TensorList)
+        mu,beta = self.get_settings(var.params, 'mu','beta', cls=NumberList)
 
-        hvp_ = as_tensorlist(hvp_)
-        update.add_(prev_update - hvp_*mu)
+        update.add_(prev_update - Hvp*mu)
         prev_update.set_(update * beta)
-        vars.update = update
-        return vars
+        var.update = update
+        return var
 
 
 class AdaptiveMatrixMomentum(Module):
+    """Second order momentum method.
+
+    Matrix momentum is useful for convex objectives, also for some reason it has very good generalization on elastic net logistic regression.
+
+    .. note::
+        In most cases MatrixMomentum should be the first module in the chain because it relies on autograd.
+
+    .. note::
+        This module requires the a closure passed to the optimizer step,
+        as it needs to re-evaluate the loss and gradients for calculating HVPs.
+        The closure must accept a ``backward`` argument (refer to documentation).
+
+
+    Args:
+        mu_mul (float, optional): multiplier to the estimated mu. Defaults to 1.
+        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
+        hvp_method (str, optional):
+            Determines how Hessian-vector products are evaluated.
+
+            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
+              This requires creating a graph for the gradient.
+            - ``"forward"``: Use a forward finite difference formula to
+              approximate the HVP. This requires one extra gradient evaluation.
+            - ``"central"``: Use a central finite difference formula for a
+              more accurate HVP approximation. This requires two extra
+              gradient evaluations.
+            Defaults to "autograd".
+        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
+        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
+
+    Reference:
+        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
     """
-    Mu here is estimated as ||s_k||/||y_k||.
-    """
-    def __init__(self, mu_mul:float=1, beta:float=1, eps=1e-4, hvp_mode: Literal['autograd', 'forward', 'central'] = 'forward', h=1e-3, hvp_tfm=None):
-        defaults = dict(mu_mul=mu_mul, beta=beta, hvp_mode=hvp_mode, h=h, eps=eps)
+
+    def __init__(
+        self,
+        mu_mul: float = 1,
+        beta: float = 1,
+        eps=1e-4,
+        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
+        h: float = 1e-3,
+        hvp_tfm: Chainable | None = None,
+    ):
+        defaults = dict(mu_mul=mu_mul, beta=beta, hvp_method=hvp_method, h=h, eps=eps)
         super().__init__(defaults)
 
         if hvp_tfm is not None:
             self.set_child('hvp_tfm', hvp_tfm)
 
+    def reset_for_online(self):
+        super().reset_for_online()
+        self.clear_state_keys('prev_params', 'prev_grad')
+
     @torch.no_grad
-    def step(self, vars):
-        assert vars.closure is not None
-        prev_update, prev_params, prev_grad = self.get_state('prev_update', 'prev_params', 'prev_grad', params=vars.params, cls=TensorList)
+    def update(self, var):
+        assert var.closure is not None
+        prev_update, prev_params, prev_grad = self.get_state(var.params, 'prev_update', 'prev_params', 'prev_grad', cls=TensorList)
 
-        settings = self.settings[vars.params[0]]
-        hvp_mode = settings['hvp_mode']
+        settings = self.settings[var.params[0]]
+        hvp_method = settings['hvp_method']
         h = settings['h']
         eps = settings['eps']
 
-        mu_mul, beta = self.get_settings('mu_mul','beta', params=vars.params, cls=NumberList)
-
-        if hvp_mode == 'autograd':
-            with torch.enable_grad():
-                grad = vars.get_grad(create_graph=True)
-                hvp_ = TensorList(hvp(vars.params, grads=grad, vec=prev_update, allow_unused=True, retain_graph=False)).detach_()
-
-        elif hvp_mode == 'forward':
-            vars.get_grad()
-            l, hvp_ = hvp_fd_forward(vars.closure, vars.params, vec=prev_update, g_0=vars.grad, h=h, normalize=True)
-            if vars.loss_approx is None: vars.loss_approx = l
+        mu_mul = NumberList(self.settings[p]['mu_mul'] for p in var.params)
 
-        elif hvp_mode == 'central':
-            l, hvp_ = hvp_fd_central(vars.closure, vars.params, vec=prev_update, h=h, normalize=True)
-            if vars.loss_approx is None: vars.loss_approx = l
-
-        else:
-            raise ValueError(hvp_mode)
+        Hvp, _ = self.Hvp(prev_update, at_x0=True, var=var, rgrad=None, hvp_method=hvp_method, h=h, normalize=True, retain_grad=False)
+        Hvp = [t.detach() for t in Hvp]
 
         if 'hvp_tfm' in self.children:
-            hvp_ = TensorList(apply(self.children['hvp_tfm'], hvp_, params=vars.params, grads=vars.grad, vars=vars))
+            Hvp = TensorList(apply_transform(self.children['hvp_tfm'], Hvp, params=var.params, grads=var.grad, var=var))
 
         # adaptive part
-        update = TensorList(vars.get_update())
-
-        s_k = vars.params - prev_params
-        prev_params.copy_(vars.params)
+        s_k = var.params - prev_params
+        prev_params.copy_(var.params)
 
-        assert vars.grad is not None
-        y_k = vars.grad - prev_grad
-        prev_grad.copy_(vars.grad)
+        if hvp_method != 'central': assert var.grad is not None
+        grad = var.get_grad()
+        y_k = grad - prev_grad
+        prev_grad.copy_(grad)
 
         ada_mu = (s_k.global_vector_norm() / (y_k.global_vector_norm() + eps)) * mu_mul
 
-        # matrix momentum uppdate
-        hvp_ = as_tensorlist(hvp_)
-        update.add_(prev_update - hvp_*ada_mu)
+        self.store(var.params, ['Hvp', 'ada_mu'], [Hvp, ada_mu])
+
+    @torch.no_grad
+    def apply(self, var):
+        Hvp, ada_mu = self.get_state(var.params, 'Hvp', 'ada_mu')
+        Hvp = as_tensorlist(Hvp)
+        beta = NumberList(self.settings[p]['beta'] for p in var.params)
+        update = TensorList(var.get_update())
+        prev_update = TensorList(self.state[p]['prev_update'] for p in var.params)
+
+        update.add_(prev_update - Hvp*ada_mu)
         prev_update.set_(update * beta)
-        vars.update = update
-        return vars
+        var.update = update
+        return var
 

torchzero/modules/momentum/momentum.py

@@ -3,11 +3,22 @@ from typing import Literal
 import torch
 
 from ...core import Target, Transform
-from ...utils import NumberList, TensorList
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
 from .ema import EMA
 
 
 class HeavyBall(EMA):
+    """Polyak's momentum (heavy-ball method).
+
+    Args:
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+        debiased (bool, optional): whether to debias the EMA like in Adam. Defaults to False.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, this becomes exponential moving average. Defaults to False.
+        ema_init (str, optional): initial values for the EMA, "zeros" or "update".
+        target (Target, optional): target to apply EMA to. Defaults to 'update'.
+    """
     def __init__(self, momentum:float=0.9, dampening:float=0, debiased: bool = False, lerp=False, ema_init: Literal['zeros', 'update'] = 'update', target: Target = 'update'):
         super().__init__(momentum=momentum, dampening=dampening, debiased=debiased, lerp=lerp, ema_init=ema_init, target=target)
 
@@ -30,14 +41,24 @@ def nag_(
 
 
 class NAG(Transform):
+    """Nesterov accelerated gradient method (nesterov momentum).
+
+    Args:
+        momentum (float, optional): momentum (beta). Defaults to 0.9.
+        dampening (float, optional): momentum dampening. Defaults to 0.
+        lerp (bool, optional):
+            whether to use linear interpolation, if True, this becomes similar to exponential moving average. Defaults to False.
+        target (Target, optional): target to apply EMA to. Defaults to 'update'.
+    """
     def __init__(self, momentum:float=0.9, dampening:float=0, lerp=False, target: Target = 'update'):
         defaults = dict(momentum=momentum,dampening=dampening, lerp=lerp)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        velocity = self.get_state('velocity', params=params, cls=TensorList)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        velocity = unpack_states(states, tensors, 'velocity', cls=TensorList)
         lerp = self.settings[params[0]]['lerp']
 
-        momentum,dampening = self.get_settings('momentum','dampening', params=params, cls=NumberList)
+        momentum,dampening = unpack_dicts(settings, 'momentum','dampening', cls=NumberList)
         return nag_(TensorList(tensors), velocity_=velocity,momentum=momentum,dampening=dampening,lerp=lerp)
+

torchzero/modules/ops/__init__.py

@@ -7,7 +7,7 @@ from .accumulate import (
 )
 from .binary import (
     Add,
-    BinaryOperation,
+    BinaryOperationBase,
     Clip,
     CopyMagnitude,
     CopySign,
@@ -27,37 +27,12 @@ from .binary import (
     Sub,
     Threshold,
 )
-from .debug import PrintShape, PrintUpdate
-from .misc import (
-    DivByLoss,
-    Dropout,
-    FillLoss,
-    GradientAccumulation,
-    GradSign,
-    GraftGradToUpdate,
-    GraftToGrad,
-    GraftToParams,
-    LastAbsoluteRatio,
-    LastDifference,
-    LastGradDifference,
-    LastProduct,
-    LastRatio,
-    MulByLoss,
-    Multistep,
-    NegateOnLossIncrease,
-    NoiseSign,
-    Previous,
-    Relative,
-    Sequential,
-    UpdateSign,
-    WeightDropout,
-)
 from .multi import (
     ClipModules,
     DivModules,
     GraftModules,
     LerpModules,
-    MultiOperation,
+    MultiOperationBase,
     PowModules,
     SubModules,
 )
@@ -66,13 +41,11 @@ from .reduce import (
     Mean,
     MinimumModules,
     Prod,
-    ReduceOperation,
+    ReduceOperationBase,
     Sum,
     WeightedMean,
     WeightedSum,
 )
-from .split import Split
-from .switch import Alternate, Switch
 from .unary import (
     Abs,
     CustomUnaryOperation,
@@ -97,7 +70,6 @@ from .utility import (
     Randn,
     RandomSample,
     Uniform,
-    Update,
     UpdateToNone,
     Zeros,
 )

torchzero/modules/ops/accumulate.py

@@ -1,65 +1,91 @@
-from collections import deque
-from operator import itemgetter
-from typing import Literal
-
 import torch
 
 from ...core import Target, Transform
-from ...utils import TensorList, NumberList
+from ...utils import TensorList, unpack_states
 
 class AccumulateSum(Transform):
+    """Accumulates sum of all past updates.
+
+    Args:
+        decay (float, optional): decays the accumulator. Defaults to 0.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, decay: float = 0, target: Target = 'update',):
         defaults = dict(decay=decay)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        sum = self.get_state('sum', params=params, cls=TensorList)
-        decay = self.get_settings('decay', params=params, cls=NumberList)
-        return sum.add_(tensors).lazy_mul(1-decay, clone=True)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        sum = unpack_states(states, tensors, 'sum', cls=TensorList)
+        decay = [1-s['decay'] for s in settings]
+        return sum.add_(tensors).lazy_mul(decay, clone=True)
 
 class AccumulateMean(Transform):
+    """Accumulates mean of all past updates.
+
+    Args:
+        decay (float, optional): decays the accumulator. Defaults to 0.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, decay: float = 0, target: Target = 'update',):
         defaults = dict(decay=decay)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-        mean = self.get_state('mean', params=params, cls=TensorList)
-        decay = self.get_settings('decay', params=params, cls=NumberList)
-        return mean.add_(tensors).lazy_mul(1-decay, clone=True).div_(step)
+        mean = unpack_states(states, tensors, 'mean', cls=TensorList)
+        decay = [1-s['decay'] for s in settings]
+        return mean.add_(tensors).lazy_mul(decay, clone=True).div_(step)
 
 class AccumulateProduct(Transform):
+    """Accumulates product of all past updates.
+
+    Args:
+        decay (float, optional): decays the accumulator. Defaults to 0.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, decay: float = 0, target: Target = 'update',):
         defaults = dict(decay=decay)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        prod = self.get_state('prod', params=params, cls=TensorList)
-        decay = self.get_settings('decay', params=params, cls=NumberList)
-        return prod.mul_(tensors).lazy_mul(1-decay, clone=True)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        prod = unpack_states(states, tensors, 'prod', cls=TensorList)
+        decay = [1-s['decay'] for s in settings]
+        return prod.mul_(tensors).lazy_mul(decay, clone=True)
 
 class AccumulateMaximum(Transform):
+    """Accumulates maximum of all past updates.
+
+    Args:
+        decay (float, optional): decays the accumulator. Defaults to 0.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, decay: float = 0, target: Target = 'update',):
         defaults = dict(decay=decay)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        maximum = self.get_state('maximum', params=params, cls=TensorList)
-        decay = self.get_settings('decay', params=params, cls=NumberList)
-        return maximum.maximum_(tensors).lazy_mul(1-decay, clone=True)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        maximum = unpack_states(states, tensors, 'maximum', cls=TensorList)
+        decay = [1-s['decay'] for s in settings]
+        return maximum.maximum_(tensors).lazy_mul(decay, clone=True)
 
 class AccumulateMinimum(Transform):
+    """Accumulates minimum of all past updates.
+
+    Args:
+        decay (float, optional): decays the accumulator. Defaults to 0.
+        target (Target, optional): target. Defaults to 'update'.
+    """
     def __init__(self, decay: float = 0, target: Target = 'update',):
         defaults = dict(decay=decay)
         super().__init__(defaults, uses_grad=False, target=target)
 
     @torch.no_grad
-    def transform(self, tensors, params, grads, vars):
-        minimum = self.get_state('minimum', params=params, cls=TensorList)
-        decay = self.get_settings('decay', params=params, cls=NumberList)
-        return minimum.minimum_(tensors).lazy_mul(1-decay, clone=True)
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        minimum = unpack_states(states, tensors, 'minimum', cls=TensorList)
+        decay = [1-s['decay'] for s in settings]
+        return minimum.minimum_(tensors).lazy_mul(decay, clone=True)