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

torchzero/modules/experimental/cosine.py

@@ -0,0 +1,214 @@
+"""A bunch of useless modules that I hate and that didn't work"""
+import torch
+
+from ...core import Chainable, Transform, apply_transform
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+
+
+class CosineStepSize(Transform):
+    """Adaptive step size based on cosine similarity
+
+    VERDICT: Useless. This is too unstable.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 0.95.
+        init (float, optional): initial step size. Defaults to 1.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        target_cossim (float, optional): cosine similarity needs to be above this to increase step size. Defaults to 1e-8.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before step size correction. Defaults to None.
+    """
+    def __init__(self, scale:float = 0.95, init:float=1, eps:float=1e-12, inner:Chainable | None = None):
+        defaults = dict(scale=scale, init=init, eps=eps)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, init = unpack_dicts(settings, 'scale', 'init', cls=NumberList)
+        unpack_states(states, tensors, 'alpha', init=init, cls=NumberList) # initializes alpha to init
+        eps = settings[0]['eps']
+
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(prev) / (tensors_norm * prev.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        new_alpha = []
+        for s, sc in zip(states, scale):
+            s['alpha'] *= 1 + cos_sim * sc
+            new_alpha.append(s['alpha'])
+
+        tensors.mul_(new_alpha)
+        prev.copy_(tensors)
+
+        return tensors
+
+
+
+class CosineDebounce(Transform):
+    """Debouncing when cosine similarity is less than 0.
+
+    VERDICT: Useless. This doesn't help at all.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 0.95.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before debouncing correction. Defaults to None.
+    """
+    def __init__(self, scale:float = 0.95, eps:float=1e-12, damping:float=0.95, inner:Chainable | None = None):
+        defaults = dict(scale=scale, eps=eps, damping=damping)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, damping = unpack_dicts(settings, 'scale', 'damping', cls=NumberList)
+        eps = settings[0]['eps']
+
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList).mul_(damping)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(prev) / (tensors_norm * prev.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        if cos_sim < -eps:
+            undo = prev.neg().mul_(-cos_sim * scale)
+            comb = prev.graft(tensors).add_(tensors).graft_(prev).mul_(-cos_sim*scale)
+            tensors = undo.add_(comb)
+
+        prev.copy_(tensors)
+        return tensors
+
+
+
+class CosineMomentum(Transform):
+    """Beta depends on cosine similarity. At cossim=1, beta is 0. At cossim=-1, beta is 2^power. This basically removes oscillations.
+
+    VERDICT: Useless. Worse than all other momentums.
+
+    Args:
+        scale (float, optional): cosine similarity multiplier. Defaults to 1.
+        nesterov (float, optional): whether to use nesterov momentum. Defaults to False.
+        power (float, optional): power for beta. Defaults to 1.
+        eps (float, optional): epsilon for division stability. Defaults to 1e-12.
+        inner (Chainable | None, optional):
+            inner modules applied after calculating cosine similarity and before updating exponential moving average. Defaults to None.
+    """
+    def __init__(self, scale:float = 1, nesterov: bool = False, power: float = 1, eps:float=1e-12, inner:Chainable | None = None):
+        defaults = dict(scale=scale, eps=eps, nesterov=nesterov, power=power)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        scale, power = unpack_dicts(settings, 'scale', 'power', cls=NumberList)
+        eps = settings[0]['eps']
+        nesterov = settings[0]['nesterov']
+        exp_avg = unpack_states(states, tensors, 'exp_avg', cls=TensorList)
+
+        tensors = as_tensorlist(tensors)
+
+        tensors_norm = tensors.global_vector_norm()
+        cos_sim = (tensors.dot(exp_avg) / (tensors_norm * exp_avg.global_vector_norm()).clip(min=eps)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        beta = (1 - (cos_sim*scale)) ** power
+        if nesterov:
+            exp_avg.add_(tensors.mul(beta))
+            return tensors.add_(exp_avg)
+        else:
+            exp_avg.add_(tensors.mul_(beta))
+            return exp_avg.clone()
+
+
+class AdaptiveDifference(Transform):
+    """VERDICT: Useless. Doesn't help (sort of to be expected)."""
+    def __init__(self, inner:Chainable | None = None):
+        defaults = dict()
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        prev = unpack_states(states, tensors, 'prev', init=tensors, cls=TensorList)
+
+        diff = tensors - prev.graft_(tensors)
+        prev.copy_(tensors)
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        tensors.add_(diff.graft_(tensors))
+
+        return tensors
+
+class AdaptiveDifferenceEMA(Transform):
+    """VERDICT: better than non-EMA but still useless."""
+    def __init__(self, beta=0.99, inner:Chainable | None = None):
+        defaults = dict(beta=beta)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        beta = unpack_dicts(settings, 'beta', cls=NumberList)
+        prev, diff_exp_avg = unpack_states(states, tensors, 'prev', 'diff_exp_avg', init=[tensors,torch.zeros_like], cls=TensorList)
+
+        diff = (tensors - prev.graft_(tensors)).graft_(tensors)
+        diff_exp_avg.lerp_(diff, 1-beta)
+        prev.copy_(tensors)
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        tensors.add_(diff_exp_avg.graft(tensors))
+
+        return tensors
+
+
+class ScaledAdaptiveDifference(Transform):
+    """VERDICT: Useless and doesn't help."""
+    def __init__(self, scale=0.95, damping:float=0.99, inner:Chainable | None = None):
+        defaults = dict(scale=scale, damping=damping)
+        super().__init__(defaults, uses_grad=False)
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        tensors = as_tensorlist(tensors)
+        scale, damping = unpack_dicts(settings, 'scale', 'damping', cls=NumberList)
+        prev_tensors, prev_update = unpack_states(states, tensors, 'prev', 'prev_update', init=[tensors,tensors], cls=TensorList)
+
+        cos_sim = (tensors.dot(prev_update) / (tensors.global_vector_norm() * prev_update.global_vector_norm()).clip(min=1e-10)).item()
+
+        if 'inner' in self.children:
+            tensors = as_tensorlist(apply_transform(self.children['inner'], tensors, params, grads, loss))
+
+        if cos_sim > 0:
+            tensors.add_(prev_tensors*(cos_sim*scale))
+
+        else:
+            undo = prev_tensors.neg().mul_(-cos_sim*scale)
+            comb = prev_tensors.graft(tensors).add_(tensors).graft_(prev_tensors).mul_(-cos_sim*scale)
+            tensors = undo.add_(comb).graft_((tensors-prev_tensors).mul_(damping))
+
+        diff = tensors - prev_tensors.graft_(tensors)
+        prev_tensors.copy_(tensors)
+        diff.graft_(tensors)
+        tensors.add_(diff)
+        prev_update.copy_(tensors)
+
+        return tensors

torchzero/modules/experimental/cubic_adam.py

@@ -0,0 +1,97 @@
+import torch
+
+from ...core import Transform
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+
+
+def signed_cbrt(x: TensorList) -> TensorList:
+    return x.sign() * x.abs().pow(1/3)
+
+def cubic_adam_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_sq_: TensorList,
+    exp_avg_cu_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    beta3: float | NumberList,
+    eps: float | NumberList,
+    debiased: bool,
+    step: int,
+):
+    exp_avg_.lerp_(tensors, 1-beta1)
+    exp_avg_sq_.lerp_(tensors**2, 1-beta2)
+    exp_avg_cu_.lerp_(tensors**3, 1-beta3)
+
+    if debiased:
+        m1 = exp_avg_ / (1 - beta1 ** step)
+        m2 = exp_avg_sq_ / (1 - beta2 ** step)
+        m3 = exp_avg_cu_ / (1 - beta3 ** step)
+    else:
+        m1, m2, m3 = exp_avg_, exp_avg_sq_, exp_avg_cu_
+
+    # adam minimizes ax^2 + bx
+    # we are going to minimize ax^3 + bx^2 + cx
+    A = signed_cbrt(m3)
+    B = m2.sqrt()
+    C = m1
+    discriminant = B.pow(2) - 4 * A * C
+
+    denom = 2 * A
+    root = discriminant.clamp(min=0).sqrt_()
+
+    x0 = (-B + root) / (denom + eps)
+    x1 = (-B - root) / (denom + eps)
+
+    f0 = (A/3)*x0**3 + (B/2)*x0**2 + C*x0
+    f1 = (A/3)*x1**3 + (B/2)*x1**2 + C*x1
+
+    x_star = x0.where(f0 < f1, x1)
+
+    adam = -C / (B + eps)
+    x_star = adam.where(discriminant < 0, x_star)
+
+    return x_star.mul_(-alpha)
+
+class CubicAdam(Transform):
+    """Adam which has 3rd momentum and minimizes a cubic polynomial.
+
+    VERDICT: can outperform Adam very slightly. Usually very similar performance.
+
+    .. warning::
+        Experimental.
+
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.99,
+        beta3: float = 0.99,
+        eps: float = 1e-8,
+        debiased:bool=True,
+        alpha: float = 1.,
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,beta3=beta3,eps=eps,debiased=debiased,alpha=alpha)
+        super().__init__(defaults, uses_grad=False)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,beta3,eps,alpha=unpack_dicts(settings, 'beta1','beta2','beta3','eps','alpha', cls=NumberList)
+        exp_avg, exp_avg_sq, exp_avg_cu = unpack_states(states, tensors, 'exp_avg', 'exp_avg_sq', 'exp_avg_cu', cls=TensorList)
+
+        return cubic_adam_(
+            tensors=TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_sq_=exp_avg_sq,
+            exp_avg_cu_=exp_avg_cu,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            beta3=beta3,
+            eps=eps,
+            debiased=settings[0]['debiased'],
+            step=step,
+        )

torchzero/modules/{projections → experimental}/dct.py

@@ -1,13 +1,13 @@
 from typing import Literal
 import torch
 import torch_dct
-from .projection import Projection
+from ..projections import ProjectionBase
 from ...core import Chainable
 
 def reverse_dims(t:torch.Tensor):
     return t.permute(*reversed(range(t.ndim)))
 
-class DCTProjection(Projection):
+class DCTProjection(ProjectionBase):
     # norm description copied from pytorch docstring
     """Project update into Discrete Cosine Transform space, requires `torch_dct` library.
 
@@ -34,8 +34,8 @@ class DCTProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, var, current):
-        settings = self.settings[var.params[0]]
+    def project(self, tensors, params, grads, loss, states, settings, current):
+        settings = settings[0]
         dims = settings['dims']
         norm = settings['norm']
 
@@ -54,18 +54,18 @@ class DCTProjection(Projection):
         return projected
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
-        settings = self.settings[var.params[0]]
+    def unproject(self, projected_tensors, params, grads, loss, projected_states, projected_settings, current):
+        settings = projected_settings[0]
         dims = settings['dims']
         norm = settings['norm']
 
         unprojected = []
-        for u in tensors:
-            dim = min(u.ndim, dims)
+        for t in projected_tensors:
+            dim = min(t.ndim, dims)
 
-            if dim == 1: idct = torch_dct.idct(u, norm = norm)
-            elif dim == 2: idct = torch_dct.idct_2d(u, norm=norm)
-            elif dim == 3: idct = torch_dct.idct_3d(u, norm=norm)
+            if dim == 1: idct = torch_dct.idct(t, norm = norm)
+            elif dim == 2: idct = torch_dct.idct_2d(t, norm=norm)
+            elif dim == 3: idct = torch_dct.idct_3d(t, norm=norm)
             else: raise ValueError(f"Unsupported number of dimensions {dim}")
 
             unprojected.append(reverse_dims(idct))

torchzero/modules/experimental/eigendescent.py

@@ -23,7 +23,10 @@ def _cosine_similarity(x, y):
 
 class EigenDescent(Module):
     """
-    Uses eigenvectors corresponding to certain eigenvalues. Please note that this is experimental and isn't guaranteed to work.
+    Uses eigenvectors corresponding to certain eigenvalues. For now they are just extracted from hessian.
+
+    .. warning::
+        Experimental.
 
     Args:
         mode (str, optional):

torchzero/modules/experimental/etf.py

@@ -4,13 +4,17 @@ import warnings
 import torch
 
 from ...core import Module
-from ...utils import vec_to_tensors, vec_to_tensors_
+from ...utils import vec_to_tensors, vec_to_tensors_, as_tensorlist
 
 
 class ExponentialTrajectoryFit(Module):
-    """A method. Please note that this is experimental and isn't guaranteed to work."""
-    def __init__(self, step_size=1e-3):
-        defaults = dict(step_size = step_size)
+    """A method.
+
+    .. warning::
+        Experimental.
+    """
+    def __init__(self, step_size=1e-2, adaptive:bool=True):
+        defaults = dict(step_size = step_size,adaptive=adaptive)
         super().__init__(defaults)
 
     @torch.no_grad
@@ -18,11 +22,17 @@ class ExponentialTrajectoryFit(Module):
         closure = var.closure
         assert closure is not None
         step_size = self.settings[var.params[0]]['step_size']
+        adaptive = self.settings[var.params[0]]['adaptive']
+
 
         # 1. perform 3 GD steps to obtain 4 points
         points = [torch.cat([p.view(-1) for p in var.params])]
         for i in range(3):
-            if i == 0: grad = var.get_grad()
+            if i == 0:
+                grad = var.get_grad()
+                if adaptive:
+                    step_size /= as_tensorlist(grad).abs().global_mean().clip(min=1e-4)
+
             else:
                 with torch.enable_grad(): closure()
                 grad = [cast(torch.Tensor, p.grad) for p in var.params]
@@ -67,9 +77,14 @@ class ExponentialTrajectoryFit(Module):
 
 
 class ExponentialTrajectoryFitV2(Module):
-    """Should be better than one above, except it isn't. Please note that this is experimental and isn't guaranteed to work."""
-    def __init__(self, step_size=1e-3, num_steps: int= 4):
-        defaults = dict(step_size = step_size, num_steps=num_steps)
+    """Should be better than one above, except it isn't.
+
+    .. warning::
+        Experimental.
+
+    """
+    def __init__(self, step_size=1e-3, num_steps: int= 4, adaptive:bool=True):
+        defaults = dict(step_size = step_size, num_steps=num_steps, adaptive=adaptive)
         super().__init__(defaults)
 
     @torch.no_grad
@@ -78,9 +93,13 @@ class ExponentialTrajectoryFitV2(Module):
         assert closure is not None
         step_size = self.settings[var.params[0]]['step_size']
         num_steps = self.settings[var.params[0]]['num_steps']
+        adaptive = self.settings[var.params[0]]['adaptive']
 
         # 1. perform 3 GD steps to obtain 4 points (or more)
         grad = var.get_grad()
+        if adaptive:
+            step_size /= as_tensorlist(grad).abs().global_mean().clip(min=1e-4)
+
         points = [torch.cat([p.view(-1) for p in var.params])]
         point_grads = [torch.cat([g.view(-1) for g in grad])]
 
@@ -132,7 +151,11 @@ def _fit_exponential(y0, y1, y2):
     return A, B, r
 
 class PointwiseExponential(Module):
-    """A stupid method (for my youtube channel). Please note that this is experimental and isn't guaranteed to work."""
+    """A stupid method (for my youtube channel).
+
+    .. warning::
+        Experimental.
+    """
     def __init__(self, step_size: float = 1e-3, reg: float = 1e-2, steps = 10000):
         defaults = dict(reg=reg, steps=steps, step_size=step_size)
         super().__init__(defaults)

torchzero/modules/experimental/exp_adam.py

@@ -0,0 +1,113 @@
+from operator import itemgetter
+from functools import partial
+import math
+import torch
+
+from ...core import Module, Target, Transform, apply_transform, Chainable
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ..functional import (
+    debias, debiased_step_size,
+    ema_,
+    sqrt_ema_sq_,
+)
+from ..step_size.lr import lazy_lr
+from ..momentum.experimental import sqrt_nag_ema_sq_
+from ..momentum.momentum import nag_
+
+
+def exp_adam_(
+    tensors: TensorList,
+    exp_avg_: TensorList,
+    exp_avg_exp_: TensorList,
+    alpha: float | NumberList,
+    beta1: float | NumberList,
+    beta2: float | NumberList,
+    eps: float | NumberList,
+    step: int,
+    pow: float = 2,
+    debiased: bool = True,
+    max_exp_avg_exp_: TensorList | None = None,
+
+    # inner args
+    inner: Module | None = None,
+    params: list[torch.Tensor] | None = None,
+    grads: list[torch.Tensor] | None = None,
+):
+    """Returns new tensors."""
+    tensors_exp = tensors.abs().clip_(max=math.log(torch.finfo(tensors[0].dtype).max) / 2).exp_()
+    exp_avg_exp_.lerp_(tensors_exp, 1-beta2)
+
+    if max_exp_avg_exp_ is not None:
+        max_exp_avg_exp_.maximum_(exp_avg_exp_)
+        exp_avg_exp_ = max_exp_avg_exp_
+
+    if inner is not None:
+        assert params is not None
+        tensors = TensorList(apply_transform(inner, tensors, params=params, grads=grads))
+
+    exp_avg_ = ema_(tensors, exp_avg_=exp_avg_, beta=beta1, dampening=0,lerp=True)
+    if debiased: alpha = debiased_step_size(step, beta1=beta1, beta2=beta2, pow=pow, alpha=alpha)
+    return (exp_avg_.lazy_mul(alpha) / exp_avg_exp_.log().add_(eps))
+
+class ExpAdam(Transform):
+    """Adam but uses abs exp and log instead of square and sqrt.
+    The gradient will be clipped to half the maximum value representable by its dtype (around 50 for float32)
+
+    Args:
+        beta1 (float, optional): momentum. Defaults to 0.9.
+        beta2 (float, optional): second momentum. Defaults to 0.999.
+        eps (float, optional): epsilon. Defaults to 1e-8.
+        alpha (float, optional): learning rate. Defaults to 1.
+        amsgrad (bool, optional): Whether to divide by maximum of EMA of gradient squares instead. Defaults to False.
+        pow (float, optional): power used in second momentum power and root. Defaults to 2.
+        debiased (bool, optional): whether to apply debiasing to momentums based on current step. Defaults to True.
+    """
+    def __init__(
+        self,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        eps: float = 1e-8,
+        amsgrad: bool = False,
+        alpha: float = 1.,
+        pow: float = 2,
+        debiased: bool = True,
+        inner: Chainable | None = None
+    ):
+        defaults=dict(beta1=beta1,beta2=beta2,eps=eps,alpha=alpha,amsgrad=amsgrad,pow=pow,debiased=debiased)
+        super().__init__(defaults, uses_grad=False)
+
+        if inner is not None: self.set_child('inner', inner)
+
+    @torch.no_grad
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+
+        beta1,beta2,eps,alpha=unpack_dicts(settings, 'beta1','beta2','eps','alpha', cls=NumberList)
+        amsgrad,pow,debiased = itemgetter('amsgrad','pow','debiased')(settings[0])
+
+        if amsgrad:
+            exp_avg, exp_avg_exp, max_exp_avg_exp = unpack_states(states, tensors, 'exp_avg', 'exp_avg_exp', 'max_exp_avg_exp', cls=TensorList)
+        else:
+            exp_avg, exp_avg_exp = unpack_states(states, tensors, 'exp_avg', 'exp_avg_exp', cls=TensorList)
+            max_exp_avg_exp = None
+
+
+        return exp_adam_(
+            tensors=TensorList(tensors),
+            exp_avg_=exp_avg,
+            exp_avg_exp_=exp_avg_exp,
+            alpha=alpha,
+            beta1=beta1,
+            beta2=beta2,
+            eps=eps,
+            step=step,
+            pow=pow,
+            debiased=debiased,
+            max_exp_avg_exp_=max_exp_avg_exp,
+
+            # inner args
+            inner=self.children.get("inner", None),
+            params=params,
+            grads=grads,
+
+        )