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

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,42 @@ _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:
+
+    ```python
+    fdm = tz.Modular(model.parameters(), tz.m.FDM(), tz.m.LR(1e-2))
+    ```
+
+    Any gradient-based method can use FDM-estimated gradients.
+    ```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
 
@@ -112,7 +138,7 @@ class FDM(GradApproximator):
             h = settings['h']
             fd_fn = _FD_FUNCS[settings['formula']]
 
-            p_flat = p.view(-1); g_flat = g.view(-1)
+            p_flat = p.ravel(); g_flat = g.ravel()
             for i in range(len(p_flat)):
                 loss, loss_approx, d = fd_fn(closure=closure, param=p_flat, idx=i, h=h, v_0=loss)
                 g_flat[i] = d

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
 
@@ -57,7 +67,9 @@ class ForwardGradient(RandomizedFDM):
         grad = None
         for i in range(n_samples):
             prt = perturbations[i]
-            if prt[0] is None: prt = params.sample_like(distribution=distribution, generator=generator)
+            if prt[0] is None:
+                prt = params.sample_like(distribution=distribution, variance=1, generator=generator)
+
             else: prt = TensorList(prt)
 
             if jvp_method == 'autograd':

torchzero/modules/grad_approximation/grad_approximator.py

@@ -14,28 +14,69 @@ 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.
+    ```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]:
-        """Returns a tuple: (grad, loss, loss_approx), make sure this resets parameters to their original values!"""
+    def approximate(self, closure: Callable, params: list[torch.Tensor], loss: torch.Tensor | None) -> tuple[Iterable[torch.Tensor], torch.Tensor | None, torch.Tensor | 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:
+    def pre_step(self, var: Var) -> None:
         """This runs once before each step, whereas `approximate` may run multiple times per step if further modules
         evaluate gradients at multiple points. This is useful for example to pre-generate new random perturbations."""
-        return var
 
     @torch.no_grad
     def step(self, var):
-        ret = self.pre_step(var)
-        if isinstance(ret, Var): var = ret
+        self.pre_step(var)
 
         if var.closure is None: raise RuntimeError("Gradient approximation requires closure")
         params, closure, loss = var.params, var.closure, var.loss
@@ -45,9 +86,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 +96,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 +104,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', 'bspsa4']