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

torchzero/modules/{projections/structural.py → experimental/structural_projections.py}

@@ -6,35 +6,18 @@ import torch
 from ...core import Chainable
 from ...utils import vec_to_tensors, TensorList
 from ..optimizers.shampoo import _merge_small_dims
-from .projection import Projection
+from ..projections import ProjectionBase
 
 
-class VectorProjection(Projection):
-    """
-    flattens and concatenates all parameters into a vector
-    """
-    def __init__(self, modules: Chainable, project_update=True, project_params=False, project_grad=False):
-        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
-    @torch.no_grad
-    def project(self, tensors, var, current):
-        return [torch.cat([u.view(-1) for u in tensors], dim=-1)]
-
-    @torch.no_grad
-    def unproject(self, tensors, var, current):
-        return vec_to_tensors(vec=tensors[0], reference=var.params)
-
-
-
-class TensorizeProjection(Projection):
+class TensorizeProjection(ProjectionBase):
     """flattens and concatenates all parameters into a vector and then reshapes it into a tensor"""
     def __init__(self, modules: Chainable, max_side: int, project_update=True, project_params=False, project_grad=False):
         defaults = dict(max_side=max_side)
         super().__init__(modules, defaults=defaults, project_update=project_update, project_params=project_params, project_grad=project_grad)
 
     @torch.no_grad
-    def project(self, tensors, var, current):
-        params = var.params
+    def project(self, tensors, params, grads, loss, states, settings, current):
         max_side = self.settings[params[0]]['max_side']
         num_elems = sum(t.numel() for t in tensors)
 
@@ -60,23 +43,23 @@ class TensorizeProjection(Projection):
         return [vec.view(dims)]
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
         remainder = self.global_state['remainder']
         # warnings.warn(f'{tensors[0].shape = }')
-        vec = tensors[0].view(-1)
+        vec = projected_tensors[0].view(-1)
         if remainder > 0: vec = vec[:-remainder]
-        return vec_to_tensors(vec, var.params)
+        return vec_to_tensors(vec, params)
 
-class BlockPartition(Projection):
+class BlockPartition(ProjectionBase):
     """splits parameters into blocks (for now flatttens them and chunks)"""
     def __init__(self, modules: Chainable, max_size: int, batched: bool = False, project_update=True, project_params=False, project_grad=False):
         defaults = dict(max_size=max_size, batched=batched)
         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):
+    def project(self, tensors, params, grads, loss, states, settings, current):
         partitioned = []
-        for p,t in zip(var.params, tensors):
+        for p,t in zip(params, tensors):
             settings = self.settings[p]
             max_size = settings['max_size']
             n = t.numel()
@@ -101,10 +84,10 @@ class BlockPartition(Projection):
         return partitioned
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
-        ti = iter(tensors)
+    def unproject(self, projected_tensors, params, grads, loss, states, settings, current):
+        ti = iter(projected_tensors)
         unprojected = []
-        for p in var.params:
+        for p in params:
             settings = self.settings[p]
             n = p.numel()
 
@@ -124,28 +107,3 @@ class BlockPartition(Projection):
 
         return unprojected
 
-
-class TensorNormsProjection(Projection):
-    def __init__(self, modules: Chainable, project_update=True, project_params=False, project_grad=False):
-        super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad)
-
-    @torch.no_grad
-    def project(self, tensors, var, current):
-        orig = self.get_state(var.params, f'{current}_orig')
-        torch._foreach_copy_(orig, tensors)
-
-        norms = torch._foreach_norm(tensors)
-        self.get_state(var.params, f'{current}_orig_norms', cls=TensorList).set_(norms)
-
-        return [torch.stack(norms)]
-
-    @torch.no_grad
-    def unproject(self, tensors, var, current):
-        orig = self.get_state(var.params, f'{current}_orig')
-        orig_norms = torch.stack(self.get_state(var.params, f'{current}_orig_norms'))
-        target_norms = tensors[0]
-
-        orig_norms = torch.where(orig_norms == 0, 1, orig_norms)
-
-        torch._foreach_mul_(orig, (target_norms/orig_norms).detach().cpu().tolist())
-        return orig

torchzero/modules/experimental/subspace_preconditioners.py

@@ -38,14 +38,19 @@ def apply_subspace_preconditioner(
     return basis @ update_projected # d
 
 class RandomSubspacePreconditioning(Transform):
-    """Whitens in random slowly changing subspace. Please note that this is experimental and isn't guaranteed to work."""
+    """Whitens in random slowly changing subspace.
+
+    .. warning::
+        Experimental and this is a barebones implementation.
+
+    """
     def __init__(self, k: int, beta: float | None = 0.99, basis_beta: float | None = 0.99, inner: Chainable | None = None):
         defaults = dict(k=k, beta=beta, basis_beta=basis_beta)
         super().__init__(defaults, uses_grad=False)
 
         if inner is not None: self.set_child('inner', inner)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         settings = settings[0]
         g = torch.cat([t.view(-1) for t in tensors])
         k = settings['k']
@@ -79,7 +84,9 @@ class RandomSubspacePreconditioning(Transform):
 
 class HistorySubspacePreconditioning(Transform):
     """Whitens in subspace spanned by history of gradient differences.
-    Please note that this is experimental and isn't guaranteed to work.
+
+    .. warning::
+        Experimental and this is a barebones implementation.
 
     Args:
         beta - for preconditioner itself in the basis.
@@ -91,7 +98,7 @@ class HistorySubspacePreconditioning(Transform):
 
         if inner is not None: self.set_child('inner', inner)
 
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         settings = settings[0]
 
         g = torch.cat([t.view(-1) for t in tensors])

torchzero/modules/experimental/{tada.py → tensor_adagrad.py}

@@ -6,17 +6,21 @@ from ...core import Chainable, TensorwiseTransform
 from ...utils.linalg import matrix_power_eigh
 
 
-class TAda(TensorwiseTransform):
-    """3rd order whitening (maybe normalizes skewness). Please note that this is experimental and isn't guaranteed to work."""
+class TensorAdagrad(TensorwiseTransform):
+    """3rd order whitening (maybe normalizes skewness, but don't quote me on it).
+
+    .. warning::
+        Experimental.
+    """
     def __init__(self, history_size: int = 100, reg: float = 1e-8, update_freq: int = 1, concat_params: bool = True, inner: Chainable | None = None):
         defaults = dict(history_size=history_size, reg=reg)
         super().__init__(defaults, uses_grad=False, update_freq=update_freq, inner=inner, concat_params=concat_params)
 
     @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, settings):
-        reg = settings['reg']
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
+        reg = setting['reg']
         if 'history' not in state:
-            state['history'] = deque(maxlen=settings['history_size'])
+            state['history'] = deque(maxlen=setting['history_size'])
 
         g = tensor.view(-1)
         history = state['history']
@@ -32,7 +36,7 @@ class TAda(TensorwiseTransform):
         state['outer'] = outer.add_(I)
 
     @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, settings):
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
         outer = state['outer']
         P = matrix_power_eigh(outer, -1/2)
         return (P @ tensor.ravel()).view_as(tensor)

torchzero/modules/functional.py

@@ -7,8 +7,9 @@ storage is always indicated in the docstring.
 
 Additional functional variants are present in most module files, e.g. `adam_`, `rmsprop_`, `lion_`, etc.
 """
-
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
+from typing import overload
+import torch
 
 from ..utils import NumberList, TensorList
 
@@ -206,4 +207,13 @@ def sqrt_centered_ema_sq_(
         ema_sq_fn=lambda *a, **kw: centered_ema_sq_(*a, **kw, exp_avg_=exp_avg_)
     )
 
+@overload
+def safe_scaling_(tensors_: torch.Tensor) -> torch.Tensor: ...
+@overload
+def safe_scaling_(tensors_: TensorList) -> TensorList: ...
+def safe_scaling_(tensors_: torch.Tensor | TensorList):
+    if isinstance(tensors_, torch.Tensor): scale = 1 / tensors_.abs().sum()
+    else: scale = 1 / tensors_.abs().global_sum()
+    scale = scale.clip(min=torch.finfo(tensors_[0].dtype).eps, max=1)
+    return tensors_.mul_(scale)
 

torchzero/modules/grad_approximation/fdm.py

@@ -77,8 +77,11 @@ def _central4(closure: Callable[..., float], param:torch.Tensor, idx: int, h, v_
     return v_0, v_plus1, (v_minus2 - 8*v_minus1 + 8*v_plus1 - v_plus2) / (12 * h)
 
 _FD_FUNCS = {
+    "forward": _forward2,
     "forward2": _forward2,
+    "backward": _backward2,
     "backward2": _backward2,
+    "central": _central2,
     "central2": _central2,
     "central3": _central2, # they are the same
     "forward3": _forward3,
@@ -88,19 +91,43 @@ _FD_FUNCS = {
 
 
 class FDM(GradApproximator):
-    """Approximate gradients via finite difference method
+    """Approximate gradients via finite difference method.
+
+    .. note::
+        This module is a gradient approximator. It modifies the closure to evaluate the estimated gradients,
+        and further closure-based modules will use the modified closure. All modules after this will use estimated gradients.
 
     Args:
         h (float, optional): magnitude of parameter perturbation. Defaults to 1e-3.
         formula (_FD_Formula, optional): finite difference formula. Defaults to 'central2'.
         target (GradTarget, optional): what to set on var. Defaults to 'closure'.
+
+    Examples:
+        plain FDM:
+
+        .. code-block:: python
+
+            fdm = tz.Modular(model.parameters(), tz.m.FDM(), tz.m.LR(1e-2))
+
+        Any gradient-based method can use FDM-estimated gradients seamlessly.
+
+        .. code-block:: python
+
+            fdm_ncg = tz.Modular(
+                model.parameters(),
+                tz.m.FDM(),
+                # set hvp_method to "forward" so that it
+                # uses gradient difference instead of autograd
+                tz.m.NewtonCG(hvp_method="forward"),
+                tz.m.Backtracking()
+            )
     """
-    def __init__(self, h: float=1e-3, formula: _FD_Formula = 'central2', target: GradTarget = 'closure'):
+    def __init__(self, h: float=1e-3, formula: _FD_Formula = 'central', target: GradTarget = 'closure'):
         defaults = dict(h=h, formula=formula)
         super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def approximate(self, closure, params, loss, var):
+    def approximate(self, closure, params, loss):
         grads = []
         loss_approx = None
 

torchzero/modules/grad_approximation/forward_gradient.py

@@ -4,14 +4,21 @@ from typing import Any, Literal
 
 import torch
 
-from ...utils import Distributions, NumberList, TensorList, generic_eq
+from ...utils import Distributions, NumberList, TensorList
 from ...utils.derivatives import jvp, jvp_fd_central, jvp_fd_forward
 from .grad_approximator import GradApproximator, GradTarget
 from .rfdm import RandomizedFDM
 
 
 class ForwardGradient(RandomizedFDM):
-    """Forward gradient method, same as randomized finite difference but directional derivative is estimated via autograd (as jacobian vector product)
+    """Forward gradient method.
+
+    This method samples one or more directional derivatives evaluated via autograd jacobian-vector products. This is very similar to randomized finite difference.
+
+    .. note::
+        This module is a gradient approximator. It modifies the closure to evaluate the estimated gradients,
+        and further closure-based modules will use the modified closure. All modules after this will use estimated gradients.
+
 
     Args:
         n_samples (int, optional): number of random gradient samples. Defaults to 1.
@@ -24,6 +31,9 @@ class ForwardGradient(RandomizedFDM):
             how to calculate jacobian vector product, note that with `forward` and 'central' this is equivalent to randomized finite difference. Defaults to 'autograd'.
         h (float, optional): finite difference step size of jvp_method is set to `forward` or `central`. Defaults to 1e-3.
         target (GradTarget, optional): what to set on var. Defaults to "closure".
+
+    References:
+        Baydin, A. G., Pearlmutter, B. A., Syme, D., Wood, F., & Torr, P. (2022). Gradients without backpropagation. arXiv preprint arXiv:2202.08587.
     """
     PRE_MULTIPLY_BY_H = False
     def __init__(
@@ -41,7 +51,7 @@ class ForwardGradient(RandomizedFDM):
         self.defaults['jvp_method'] = jvp_method
 
     @torch.no_grad
-    def approximate(self, closure, params, loss, var):
+    def approximate(self, closure, params, loss):
         params = TensorList(params)
         loss_approx = None
 

torchzero/modules/grad_approximation/grad_approximator.py

@@ -14,17 +14,62 @@ class GradApproximator(Module, ABC):
     """Base class for gradient approximations.
     This is an abstract class, to use it, subclass it and override `approximate`.
 
+    GradientApproximator modifies the closure to evaluate the estimated gradients,
+    and further closure-based modules will use the modified closure.
+
     Args:
         defaults (dict[str, Any] | None, optional): dict with defaults. Defaults to None.
         target (str, optional):
             whether to set `var.grad`, `var.update` or 'var.closure`. Defaults to 'closure'.
-    """
+
+    Example:
+
+        Basic SPSA method implementation.
+
+        .. code-block:: python
+
+            class SPSA(GradApproximator):
+                def __init__(self, h=1e-3):
+                    defaults = dict(h=h)
+                    super().__init__(defaults)
+
+                @torch.no_grad
+                def approximate(self, closure, params, loss):
+                    perturbation = [rademacher_like(p) * self.settings[p]['h'] for p in params]
+
+                    # evaluate params + perturbation
+                    torch._foreach_add_(params, perturbation)
+                    loss_plus = closure(False)
+
+                    # evaluate params - perturbation
+                    torch._foreach_sub_(params, perturbation)
+                    torch._foreach_sub_(params, perturbation)
+                    loss_minus = closure(False)
+
+                    # restore original params
+                    torch._foreach_add_(params, perturbation)
+
+                    # calculate SPSA gradients
+                    spsa_grads = []
+                    for p, pert in zip(params, perturbation):
+                        settings = self.settings[p]
+                        h = settings['h']
+                        d = (loss_plus - loss_minus) / (2*(h**2))
+                        spsa_grads.append(pert * d)
+
+                    # returns tuple: (grads, loss, loss_approx)
+                    # loss must be with initial parameters
+                    # since we only evaluated loss with perturbed parameters
+                    # we only have loss_approx
+                    return spsa_grads, None, loss_plus
+
+            """
     def __init__(self, defaults: dict[str, Any] | None = None, target: GradTarget = 'closure'):
         super().__init__(defaults)
         self._target: GradTarget = target
 
     @abstractmethod
-    def approximate(self, closure: Callable, params: list[torch.Tensor], loss: _Scalar | None, var: Var) -> tuple[Iterable[torch.Tensor], _Scalar | None, _Scalar | None]:
+    def approximate(self, closure: Callable, params: list[torch.Tensor], loss: _Scalar | None) -> tuple[Iterable[torch.Tensor], _Scalar | None, _Scalar | None]:
         """Returns a tuple: (grad, loss, loss_approx), make sure this resets parameters to their original values!"""
 
     def pre_step(self, var: Var) -> Var | None:
@@ -45,9 +90,9 @@ class GradApproximator(Module, ABC):
             def approx_closure(backward=True):
                 if backward:
                     # set loss to None because closure might be evaluated at different points
-                    grad, l, l_approx = self.approximate(closure=closure, params=params, loss=None, var=var)
+                    grad, l, l_approx = self.approximate(closure=closure, params=params, loss=None)
                     for p, g in zip(params, grad): p.grad = g
-                    return l if l is not None else l_approx
+                    return l if l is not None else closure(False)
                 return closure(False)
 
             var.closure = approx_closure
@@ -55,7 +100,7 @@ class GradApproximator(Module, ABC):
 
         # if var.grad is not None:
         #     warnings.warn('Using grad approximator when `var.grad` is already set.')
-        grad,loss,loss_approx = self.approximate(closure=closure, params=params, loss=loss, var=var)
+        grad,loss,loss_approx = self.approximate(closure=closure, params=params, loss=loss)
         if loss_approx is not None: var.loss_approx = loss_approx
         if loss is not None: var.loss = var.loss_approx = loss
         if self._target == 'grad': var.grad = list(grad)
@@ -63,4 +108,4 @@ class GradApproximator(Module, ABC):
         else: raise ValueError(self._target)
         return var
 
-_FD_Formula = Literal['forward2', 'backward2', 'forward3', 'backward3', 'central2', 'central4']
+_FD_Formula = Literal['forward', 'forward2', 'backward', 'backward2', 'central', 'central2', 'central3', 'forward3', 'backward3', 'central4', 'forward4', 'forward5', 'bspsa5']