torchzero 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

torchzero/modules/experimental/spsa1.py

@@ -38,15 +38,15 @@ class SPSA1(GradApproximator):
         super().__init__(defaults, target=target)
 
 
-    def pre_step(self, var):
+    def pre_step(self, objective):
 
         if self.defaults['pre_generate']:
 
-            params = TensorList(var.params)
+            params = TensorList(objective.params)
             generator = self.get_generator(params[0].device, self.defaults['seed'])
 
             n_samples = self.defaults['n_samples']
-            h = self.get_settings(var.params, 'h')
+            h = self.get_settings(objective.params, 'h')
 
             perturbations = [params.rademacher_like(generator=generator) for _ in range(n_samples)]
             torch._foreach_mul_([p for l in perturbations for p in l], [v for vv in h for v in [vv]*n_samples])

torchzero/modules/experimental/structural_projections.py

@@ -1,11 +1,8 @@
 import math
-
-import numpy as np
 import torch
 
 from ...core import Chainable
-from ...utils import vec_to_tensors, TensorList
-from ..adaptive.shampoo import _merge_small_dims
+from ...utils import vec_to_tensors
 from ..projections import ProjectionBase
 
 

torchzero/modules/grad_approximation/fdm.py

@@ -106,12 +106,12 @@ class FDM(GradApproximator):
     plain FDM:
 
     ```python
-    fdm = tz.Modular(model.parameters(), tz.m.FDM(), tz.m.LR(1e-2))
+    fdm = tz.Optimizer(model.parameters(), tz.m.FDM(), tz.m.LR(1e-2))
     ```
 
     Any gradient-based method can use FDM-estimated gradients.
     ```python
-    fdm_ncg = tz.Modular(
+    fdm_ncg = tz.Optimizer(
         model.parameters(),
         tz.m.FDM(),
         # set hvp_method to "forward" so that it

torchzero/modules/grad_approximation/forward_gradient.py

@@ -52,11 +52,11 @@ class ForwardGradient(RandomizedFDM):
         params = TensorList(params)
         loss_approx = None
 
-        settings = self.settings[params[0]]
-        n_samples = settings['n_samples']
-        jvp_method = settings['jvp_method']
-        h = settings['h']
-        distribution = settings['distribution']
+        fs = self.settings[params[0]]
+        n_samples = fs['n_samples']
+        jvp_method = fs['jvp_method']
+        h = fs['h']
+        distribution = fs['distribution']
         default = [None]*n_samples
         perturbations = list(zip(*(self.state[p].get('perturbations', default) for p in params)))
         generator = self.get_generator(params[0].device, self.defaults['seed'])
@@ -74,10 +74,10 @@ class ForwardGradient(RandomizedFDM):
                     loss, d = jvp(partial(closure, False), params=params, tangent=prt)
 
             elif jvp_method == 'forward':
-                loss, d = jvp_fd_forward(partial(closure, False), params=params, tangent=prt, v_0=loss, normalize=True, h=h)
+                loss, d = jvp_fd_forward(partial(closure, False), params=params, tangent=prt, v_0=loss, h=h)
 
             elif jvp_method == 'central':
-                loss_approx, d = jvp_fd_central(partial(closure, False), params=params, tangent=prt, normalize=True, h=h)
+                loss_approx, d = jvp_fd_central(partial(closure, False), params=params, tangent=prt, h=h)
 
             else: raise ValueError(jvp_method)
 

torchzero/modules/grad_approximation/grad_approximator.py

@@ -5,7 +5,7 @@ from typing import Any, Literal
 
 import torch
 
-from ...core import Module, Var
+from ...core import Module, Objective
 
 GradTarget = Literal['update', 'grad', 'closure']
 _Scalar = torch.Tensor | float
@@ -62,24 +62,25 @@ class GradApproximator(Module, ABC):
             return spsa_grads, None, loss_plus
     ```
     """
-    def __init__(self, defaults: dict[str, Any] | None = None, target: GradTarget = 'closure'):
+    def __init__(self, defaults: dict[str, Any] | None = None, return_approx_loss:bool=False, target: GradTarget = 'closure'):
         super().__init__(defaults)
         self._target: GradTarget = target
+        self._return_approx_loss = return_approx_loss
 
     @abstractmethod
     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) -> None:
+    def pre_step(self, objective: Objective) -> 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."""
 
     @torch.no_grad
-    def step(self, var):
-        self.pre_step(var)
+    def update(self, objective):
+        self.pre_step(objective)
 
-        if var.closure is None: raise RuntimeError("Gradient approximation requires closure")
-        params, closure, loss = var.params, var.closure, var.loss
+        if objective.closure is None: raise RuntimeError("Gradient approximation requires closure")
+        params, closure, loss = objective.params, objective.closure, objective.loss
 
         if self._target == 'closure':
 
@@ -88,20 +89,26 @@ class GradApproximator(Module, ABC):
                     # set loss to None because closure might be evaluated at different points
                     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 closure(False)
+                    if l is not None: return l
+                    if self._return_approx_loss and l_approx is not None: return l_approx
+                    return closure(False)
+
                 return closure(False)
 
-            var.closure = approx_closure
-            return var
+            objective.closure = approx_closure
+            return
 
         # 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)
-        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)
-        elif self._target == 'update': var.update = list(grad)
+        grad, loss, loss_approx = self.approximate(closure=closure, params=params, loss=loss)
+        if loss_approx is not None: objective.loss_approx = loss_approx
+        if loss is not None: objective.loss = objective.loss_approx = loss
+        if self._target == 'grad': objective.grads = list(grad)
+        elif self._target == 'update': objective.updates = list(grad)
         else: raise ValueError(self._target)
-        return var
+        return
+
+    def apply(self, objective):
+        return objective
 
 _FD_Formula = Literal['forward', 'forward2', 'backward', 'backward2', 'central', 'central2', 'central3', 'forward3', 'backward3', 'central4', 'forward4', 'forward5', 'bspsa4']

torchzero/modules/grad_approximation/rfdm.py

@@ -174,9 +174,9 @@ class RandomizedFDM(GradApproximator):
 
     SPSA is randomized FDM with rademacher distribution and central formula.
     ```py
-    spsa = tz.Modular(
+    spsa = tz.Optimizer(
         model.parameters(),
-        tz.m.RandomizedFDM(formula="central", distribution="rademacher"),
+        tz.m.RandomizedFDM(formula="fd_central", distribution="rademacher"),
         tz.m.LR(1e-2)
     )
     ```
@@ -185,9 +185,9 @@ class RandomizedFDM(GradApproximator):
 
     RDSA is randomized FDM with usually gaussian distribution and central formula.
     ```
-    rdsa = tz.Modular(
+    rdsa = tz.Optimizer(
         model.parameters(),
-        tz.m.RandomizedFDM(formula="central", distribution="gaussian"),
+        tz.m.RandomizedFDM(formula="fd_central", distribution="gaussian"),
         tz.m.LR(1e-2)
     )
     ```
@@ -196,7 +196,7 @@ class RandomizedFDM(GradApproximator):
 
     GS uses many gaussian samples with possibly a larger finite difference step size.
     ```
-    gs = tz.Modular(
+    gs = tz.Optimizer(
         model.parameters(),
         tz.m.RandomizedFDM(n_samples=100, distribution="gaussian", formula="forward2", h=1e-1),
         tz.m.NewtonCG(hvp_method="forward"),
@@ -208,7 +208,7 @@ class RandomizedFDM(GradApproximator):
 
     Momentum might help by reducing the variance of the estimated gradients.
     ```
-    momentum_spsa = tz.Modular(
+    momentum_spsa = tz.Optimizer(
         model.parameters(),
         tz.m.RandomizedFDM(),
         tz.m.HeavyBall(0.9),
@@ -223,23 +223,24 @@ class RandomizedFDM(GradApproximator):
         n_samples: int = 1,
         formula: _FD_Formula = "central",
         distribution: Distributions = "rademacher",
-        pre_generate = True,
+        pre_generate: bool = True,
+        return_approx_loss: bool = False,
         seed: int | None | torch.Generator = None,
         target: GradTarget = "closure",
     ):
         defaults = dict(h=h, formula=formula, n_samples=n_samples, distribution=distribution, pre_generate=pre_generate, seed=seed)
-        super().__init__(defaults, target=target)
+        super().__init__(defaults, return_approx_loss=return_approx_loss, target=target)
 
 
-    def pre_step(self, var):
-        h = self.get_settings(var.params, 'h')
+    def pre_step(self, objective):
+        h = self.get_settings(objective.params, 'h')
         pre_generate = self.defaults['pre_generate']
 
         if pre_generate:
             n_samples = self.defaults['n_samples']
             distribution = self.defaults['distribution']
 
-            params = TensorList(var.params)
+            params = TensorList(objective.params)
             generator = self.get_generator(params[0].device, self.defaults['seed'])
             perturbations = [params.sample_like(distribution=distribution, variance=1, generator=generator) for _ in range(n_samples)]
 
@@ -346,11 +347,12 @@ class RDSA(RandomizedFDM):
         n_samples: int = 1,
         formula: _FD_Formula = "central2",
         distribution: Distributions = "gaussian",
-        pre_generate = True,
+        pre_generate: bool = True,
+        return_approx_loss: bool = False,
         target: GradTarget = "closure",
         seed: int | None | torch.Generator = None,
     ):
-        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,pre_generate=pre_generate,target=target,seed=seed)
+        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,pre_generate=pre_generate,target=target,seed=seed, return_approx_loss=return_approx_loss)
 
 class GaussianSmoothing(RandomizedFDM):
     """
@@ -380,11 +382,12 @@ class GaussianSmoothing(RandomizedFDM):
         n_samples: int = 100,
         formula: _FD_Formula = "forward2",
         distribution: Distributions = "gaussian",
-        pre_generate = True,
+        pre_generate: bool = True,
+        return_approx_loss: bool = False,
         target: GradTarget = "closure",
         seed: int | None | torch.Generator = None,
     ):
-        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,pre_generate=pre_generate,target=target,seed=seed)
+        super().__init__(h=h, n_samples=n_samples,formula=formula,distribution=distribution,pre_generate=pre_generate,target=target,seed=seed, return_approx_loss=return_approx_loss)
 
 class MeZO(GradApproximator):
     """Gradient approximation via memory-efficient zeroth order optimizer (MeZO) - https://arxiv.org/abs/2305.17333.
@@ -406,10 +409,10 @@ class MeZO(GradApproximator):
     """
 
     def __init__(self, h: float=1e-3, n_samples: int = 1, formula: _FD_Formula = 'central2',
-                 distribution: Distributions = 'rademacher', target: GradTarget = 'closure'):
+                 distribution: Distributions = 'rademacher', return_approx_loss: bool = False, target: GradTarget = 'closure'):
 
         defaults = dict(h=h, formula=formula, n_samples=n_samples, distribution=distribution)
-        super().__init__(defaults, target=target)
+        super().__init__(defaults, return_approx_loss=return_approx_loss, target=target)
 
     def _seeded_perturbation(self, params: list[torch.Tensor], distribution, seed, h):
         prt = TensorList(params).sample_like(
@@ -419,19 +422,19 @@ class MeZO(GradApproximator):
         )
         return prt
 
-    def pre_step(self, var):
-        h = NumberList(self.settings[p]['h'] for p in var.params)
+    def pre_step(self, objective):
+        h = NumberList(self.settings[p]['h'] for p in objective.params)
 
         n_samples = self.defaults['n_samples']
         distribution = self.defaults['distribution']
 
-        step = var.current_step
+        step = objective.current_step
 
         # create functions that generate a deterministic perturbation from seed based on current step
         prt_fns = []
         for i in range(n_samples):
 
-            prt_fn = partial(self._seeded_perturbation, params=var.params, distribution=distribution, seed=1_000_000*step + i, h=h)
+            prt_fn = partial(self._seeded_perturbation, params=objective.params, distribution=distribution, seed=1_000_000*step + i, h=h)
             prt_fns.append(prt_fn)
 
         self.global_state['prt_fns'] = prt_fns

torchzero/modules/least_squares/gn.py

@@ -1,28 +1,31 @@
 import torch
-from ...core import Module
 
-from ...utils.derivatives import jacobian_wrt, flatten_jacobian
+from ...core import Chainable, Transform
+from ...linalg import linear_operator
 from ...utils import vec_to_tensors
-from ...utils.linalg import linear_operator
-class SumOfSquares(Module):
+from ...utils.derivatives import flatten_jacobian, jacobian_wrt
+
+
+class SumOfSquares(Transform):
     """Sets loss to be the sum of squares of values returned by the closure.
 
     This is meant to be used to test least squares methods against ordinary minimization methods.
 
     To use this, the closure should return a vector of values to minimize sum of squares of.
-    Please add the `backward` argument, it will always be False but it is required.
+    Please add the ``backward`` argument, it will always be False but it is required.
     """
     def __init__(self):
         super().__init__()
 
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def update_states(self, objective, states, settings):
+        closure = objective.closure
 
         if closure is not None:
+
             def sos_closure(backward=True):
                 if backward:
-                    var.zero_grad()
+                    objective.zero_grad()
                     with torch.enable_grad():
                         loss = closure(False)
                         loss = loss.pow(2).sum()
@@ -32,18 +35,19 @@ class SumOfSquares(Module):
                 loss = closure(False)
                 return loss.pow(2).sum()
 
-            var.closure = sos_closure
-
-        if var.loss is not None:
-            var.loss = var.loss.pow(2).sum()
+            objective.closure = sos_closure
 
-        if var.loss_approx is not None:
-            var.loss_approx = var.loss_approx.pow(2).sum()
+        if objective.loss is not None:
+            objective.loss = objective.loss.pow(2).sum()
 
-        return var
+        if objective.loss_approx is not None:
+            objective.loss_approx = objective.loss_approx.pow(2).sum()
 
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        return objective
 
-class GaussNewton(Module):
+class GaussNewton(Transform):
     """Gauss-newton method.
 
     To use this, the closure should return a vector of values to minimize sum of squares of.
@@ -57,6 +61,9 @@ class GaussNewton(Module):
 
     Args:
         reg (float, optional): regularization parameter. Defaults to 1e-8.
+        update_freq (int, optional):
+            frequency of computing the jacobian. When jacobian is not computed, only residuals are computed and updated.
+            Defaults to 1.
         batched (bool, optional): whether to use vmapping. Defaults to True.
 
     Examples:
@@ -68,7 +75,7 @@ class GaussNewton(Module):
         return torch.stack([(1 - x1), 100 * (x2 - x1**2)])
 
     X = torch.tensor([-1.1, 2.5], requires_grad=True)
-    opt = tz.Modular([X], tz.m.GaussNewton(), tz.m.Backtracking())
+    opt = tz.Optimizer([X], tz.m.GaussNewton(), tz.m.Backtracking())
 
     # define the closure for line search
     def closure(backward=True):
@@ -86,7 +93,7 @@ class GaussNewton(Module):
     y = torch.randn(64, 10)
 
     model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.TrustCG(tz.m.GaussNewton()),
     )
@@ -101,35 +108,62 @@ class GaussNewton(Module):
             print(f'{losses.mean() = }')
     ```
     """
-    def __init__(self, reg:float = 1e-8, batched:bool=True, ):
-        super().__init__(defaults=dict(batched=batched, reg=reg))
+    def __init__(self, reg:float = 1e-8, update_freq: int= 1, batched:bool=True, inner: Chainable | None = None):
+        defaults=dict(update_freq=update_freq,batched=batched, reg=reg)
+        super().__init__(defaults=defaults)
+        if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def update(self, var):
-        params = var.params
-        batched = self.defaults['batched']
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        params = objective.params
+        closure = objective.closure
+        batched = fs['batched']
+        update_freq = fs['update_freq']
+
+        # compute residuals
+        r = objective.loss
+        if r is None:
+            assert closure is not None
+            with torch.enable_grad():
+                r = objective.get_loss(backward=False) # n_residuals
+                assert isinstance(r, torch.Tensor)
+
+        # set sum of squares scalar loss and it's gradient to objective
+        objective.loss = r.pow(2).sum()
+
+        step = self.increment_counter("step", start=0)
+
+        if step % update_freq == 0:
+
+            # compute jacobian
+            with torch.enable_grad():
+                J_list = jacobian_wrt([r.ravel()], params, batched=batched)
+
+            J = self.global_state["J"] = flatten_jacobian(J_list) # (n_residuals, ndim)
 
-        closure = var.closure
-        assert closure is not None
+        else:
+            J = self.global_state["J"]
 
-        # gauss newton direction
-        with torch.enable_grad():
-            f = var.get_loss(backward=False) # n_out
-            assert isinstance(f, torch.Tensor)
-            G_list = jacobian_wrt([f.ravel()], params, batched=batched)
+        Jr = J.T @ r.detach() # (ndim)
 
-        var.loss = f.pow(2).sum()
+        # if there are more residuals, solve (J^T J)x = J^T r, so we need Jr
+        # otherwise solve (J J^T)z = r and set x = J^T z, so we need r
+        n_residuals, ndim = J.shape
+        if n_residuals >= ndim or "inner" in self.children:
+            self.global_state["Jr"] = Jr
 
-        G = self.global_state["G"] = flatten_jacobian(G_list) # (n_out, ndim)
-        Gtf = G.T @ f.detach() # (ndim)
-        self.global_state["Gtf"] = Gtf
-        var.grad = vec_to_tensors(Gtf, var.params)
+        else:
+            self.global_state["r"] = r
+
+        objective.grads = vec_to_tensors(Jr, objective.params)
 
         # set closure to calculate sum of squares for line searches etc
-        if var.closure is not None:
+        if closure is not None:
             def sos_closure(backward=True):
+
                 if backward:
-                    var.zero_grad()
+                    objective.zero_grad()
                     with torch.enable_grad():
                         loss = closure(False).pow(2).sum()
                         loss.backward()
@@ -138,24 +172,61 @@ class GaussNewton(Module):
                 loss = closure(False).pow(2).sum()
                 return loss
 
-            var.closure = sos_closure
+            objective.closure = sos_closure
 
     @torch.no_grad
-    def apply(self, var):
-        reg = self.defaults['reg']
+    def apply_states(self, objective, states, settings):
+        fs = settings[0]
+        reg = fs['reg']
+
+        J: torch.Tensor = self.global_state['J']
+        nresiduals, ndim = J.shape
+        if nresiduals >= ndim or "inner" in self.children:
+
+            # (J^T J)v = J^T r
+            Jr: torch.Tensor = self.global_state['Jr']
+
+            # inner step
+            if "inner" in self.children:
+
+                # var.grad is set to unflattened Jr
+                assert objective.grads is not None
+                objective = self.inner_step("inner", objective, must_exist=True)
+                Jr_list = objective.get_updates()
+                Jr = torch.cat([t.ravel() for t in Jr_list])
+
+            JtJ = J.T @ J # (ndim, ndim)
+            if reg != 0:
+                JtJ.add_(torch.eye(JtJ.size(0), device=JtJ.device, dtype=JtJ.dtype).mul_(reg))
+
+            if nresiduals >= ndim:
+                v, info = torch.linalg.solve_ex(JtJ, Jr) # pylint:disable=not-callable
+            else:
+                v = torch.linalg.lstsq(JtJ, Jr).solution # pylint:disable=not-callable
+
+            objective.updates = vec_to_tensors(v, objective.params)
+            return objective
+
+        # else:
+        # solve (J J^T)z = r and set v = J^T z
+        # we need (J^T J)v = J^T r
+        # if z is solution to (G G^T)z = r, and v = J^T z
+        # then (J^T J)v = (J^T J) (J^T z) = J^T (J J^T) z = J^T r
+        # therefore (J^T J)v = J^T r
+        # also this gives a minimum norm solution
 
-        G = self.global_state['G']
-        Gtf = self.global_state['Gtf']
+        r = self.global_state['r']
 
-        GtG = G.T @ G # (ndim, ndim)
+        JJT = J @ J.T # (nresiduals, nresiduals)
         if reg != 0:
-            GtG.add_(torch.eye(GtG.size(0), device=GtG.device, dtype=GtG.dtype).mul_(reg))
+            JJT.add_(torch.eye(JJT.size(0), device=JJT.device, dtype=JJT.dtype).mul_(reg))
 
-        v = torch.linalg.lstsq(GtG, Gtf).solution # pylint:disable=not-callable
+        z, info = torch.linalg.solve_ex(JJT, r) # pylint:disable=not-callable
+        v = J.T @ z
 
-        var.update = vec_to_tensors(v, var.params)
-        return var
+        objective.updates = vec_to_tensors(v, objective.params)
+        return objective
 
-    def get_H(self, var):
-        G = self.global_state['G']
-        return linear_operator.AtA(G)
+    def get_H(self, objective=...):
+        J = self.global_state['J']
+        return linear_operator.AtA(J)

torchzero/modules/line_search/backtracking.py

@@ -77,7 +77,7 @@ class Backtracking(LineSearchBase):
     Gradient descent with backtracking line search:
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Backtracking()
     )
@@ -85,7 +85,7 @@ class Backtracking(LineSearchBase):
 
     L-BFGS with backtracking line search:
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.LBFGS(),
         tz.m.Backtracking()
@@ -117,7 +117,7 @@ class Backtracking(LineSearchBase):
 
         # # directional derivative
         if c == 0: d = 0
-        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), var.get_update()))
+        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grads(), var.get_updates()))
 
         # scale init
         init_scale = self.global_state.get('init_scale', 1)
@@ -199,7 +199,7 @@ class AdaptiveBacktracking(LineSearchBase):
 
         # directional derivative (0 if c = 0 because it is not needed)
         if c == 0: d = 0
-        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grad(), update))
+        else: d = -sum(t.sum() for t in torch._foreach_mul(var.get_grads(), update))
 
         # scale beta
         beta = beta * self.global_state['beta_scale']