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

torchzero/modules/experimental/scipy_newton_cg.py

@@ -3,10 +3,9 @@ from typing import Literal, overload
 import torch
 from scipy.sparse.linalg import LinearOperator, gcrotmk
 
-from ...core import Chainable, Module, apply_transform
-from ...utils import NumberList, TensorList, as_tensorlist, generic_vector_norm, vec_to_tensors
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-from ...utils.linalg.solve import cg, minres
+from ...core import Chainable, Module, step
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import hvp_fd_central, hvp_fd_forward
 
 
 class ScipyNewtonCG(Module):
@@ -14,7 +13,7 @@ class ScipyNewtonCG(Module):
     def __init__(
         self,
         solver = gcrotmk,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        hvp_method: Literal["fd_forward", "fd_central", "autograd"] = "autograd",
         h: float = 1e-3,
         warm_start=False,
         inner: Chainable | None = None,
@@ -33,47 +32,47 @@ class ScipyNewtonCG(Module):
         self._kwargs = kwargs
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
+    def apply(self, objective):
+        params = TensorList(objective.params)
+        closure = objective.closure
         if closure is None: raise RuntimeError('NewtonCG requires closure')
 
-        settings = self.settings[params[0]]
-        hvp_method = settings['hvp_method']
-        solver = settings['solver']
-        h = settings['h']
-        warm_start = settings['warm_start']
+        fs = self.settings[params[0]]
+        hvp_method = fs['hvp_method']
+        solver = fs['solver']
+        h = fs['h']
+        warm_start = fs['warm_start']
 
         self._num_hvps_last_step = 0
         # ---------------------- Hessian vector product function --------------------- #
         device = params[0].device; dtype=params[0].dtype
         if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
+            grad = objective.get_grads(create_graph=True)
 
             def H_mm(x_np):
                 self._num_hvps_last_step += 1
                 x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
                 with torch.enable_grad():
-                    Hvp = TensorList(hvp(params, grad, x, retain_graph=True))
+                    Hvp = TensorList(torch.autograd.grad(grad, params, x, retain_graph=True))
                 return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
         else:
 
             with torch.enable_grad():
-                grad = var.get_grad()
+                grad = objective.get_grads()
 
             if hvp_method == 'forward':
                 def H_mm(x_np):
                     self._num_hvps_last_step += 1
                     x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
-                    Hvp = TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
+                    Hvp = TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad)[1])
                     return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
             elif hvp_method == 'central':
                 def H_mm(x_np):
                     self._num_hvps_last_step += 1
                     x = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), grad)
-                    Hvp = TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
+                    Hvp = TensorList(hvp_fd_central(closure, params, x, h=h)[1])
                     return torch.cat([t.ravel() for t in Hvp]).numpy(force=True)
 
             else:
@@ -83,10 +82,8 @@ class ScipyNewtonCG(Module):
         H = LinearOperator(shape=(ndim,ndim), matvec=H_mm, rmatvec=H_mm) # type:ignore
 
         # -------------------------------- inner step -------------------------------- #
-        b = var.get_update()
-        if 'inner' in self.children:
-            b = apply_transform(self.children['inner'], b, params=params, grads=grad, var=var)
-        b = as_tensorlist(b)
+        objective = self.inner_step("inner", objective, must_exist=False)
+        b = TensorList(objective.get_updates())
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
@@ -98,8 +95,8 @@ class ScipyNewtonCG(Module):
         if warm_start:
             self.global_state['x_prev'] = x_np
 
-        var.update = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), params)
+        objective.updates = vec_to_tensors(torch.as_tensor(x_np, device=device, dtype=dtype), params)
 
         self._num_hvps += self._num_hvps_last_step
-        return var
+        return objective
 

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/functional.py

@@ -30,7 +30,7 @@ def debiased_step_size(
     pow: float = 2,
     alpha: float | NumberList = 1,
 ):
-    """returns multiplier to step size"""
+    """returns multiplier to step size, step starts from 1"""
     if isinstance(beta1, NumberList): beta1 = beta1.fill_none(0)
     if isinstance(beta2, NumberList): beta2 = beta2.fill_none(0)
 

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

@@ -176,7 +176,7 @@ class RandomizedFDM(GradApproximator):
     ```py
     spsa = tz.Modular(
         model.parameters(),
-        tz.m.RandomizedFDM(formula="central", distribution="rademacher"),
+        tz.m.RandomizedFDM(formula="fd_central", distribution="rademacher"),
         tz.m.LR(1e-2)
     )
     ```
@@ -187,7 +187,7 @@ class RandomizedFDM(GradApproximator):
     ```
     rdsa = tz.Modular(
         model.parameters(),
-        tz.m.RandomizedFDM(formula="central", distribution="gaussian"),
+        tz.m.RandomizedFDM(formula="fd_central", distribution="gaussian"),
         tz.m.LR(1e-2)
     )
     ```
@@ -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, Module, step
+from ...linalg import linear_operator
 from ...utils import vec_to_tensors
-from ...utils.linalg import linear_operator
+from ...utils.derivatives import flatten_jacobian, jacobian_wrt
+
+
 class SumOfSquares(Module):
     """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(self, objective):
+        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,16 +35,13 @@ 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()
-
-        return var
+        if objective.loss is not None:
+            objective.loss = objective.loss.pow(2).sum()
 
+        if objective.loss_approx is not None:
+            objective.loss_approx = objective.loss_approx.pow(2).sum()
 
 class GaussNewton(Module):
     """Gauss-newton method.
@@ -101,35 +101,45 @@ class GaussNewton(Module):
             print(f'{losses.mean() = }')
     ```
     """
-    def __init__(self, reg:float = 1e-8, batched:bool=True, ):
+    def __init__(self, reg:float = 1e-8, batched:bool=True, inner: Chainable | None = None):
         super().__init__(defaults=dict(batched=batched, reg=reg))
+        if inner is not None: self.set_child('inner', inner)
 
     @torch.no_grad
-    def update(self, var):
-        params = var.params
+    def update(self, objective):
+        params = objective.params
         batched = self.defaults['batched']
 
-        closure = var.closure
+        closure = objective.closure
         assert closure is not None
 
         # 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)
+            r = objective.get_loss(backward=False) # nresiduals
+            assert isinstance(r, torch.Tensor)
+            J_list = jacobian_wrt([r.ravel()], params, batched=batched)
+
+        objective.loss = r.pow(2).sum()
+
+        J = self.global_state["J"] = flatten_jacobian(J_list) # (nresiduals, ndim)
+        Jr = J.T @ r.detach() # (ndim)
+
+        # 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
+        nresiduals, ndim = J.shape
+        if nresiduals >= ndim or "inner" in self.children:
+            self.global_state["Jr"] = Jr
 
-        var.loss = f.pow(2).sum()
+        else:
+            self.global_state["r"] = r
 
-        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)
+        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 objective.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 +148,62 @@ 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):
+    def apply(self, objective):
         reg = self.defaults['reg']
 
-        G = self.global_state['G']
-        Gtf = self.global_state['Gtf']
+        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])
+
+            JJ = J.T @ J # (ndim, ndim)
+            if reg != 0:
+                JJ.add_(torch.eye(JJ.size(0), device=JJ.device, dtype=JJ.dtype).mul_(reg))
+
+            if nresiduals >= ndim:
+                v, info = torch.linalg.solve_ex(JJ, Jr) # pylint:disable=not-callable
+            else:
+                v = torch.linalg.lstsq(JJ, 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
+            # derivation
+            # we need (J^T J)v = J^T r
+            # suppose z is solution to (G G^T)z = r, and v = J^T z
+            # if 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 with our presuppositions (J^T J)v = J^T r
+
+            # also this gives a minimum norm solution
+
+            r = self.global_state['r']
 
-        GtG = G.T @ G # (ndim, ndim)
-        if reg != 0:
-            GtG.add_(torch.eye(GtG.size(0), device=GtG.device, dtype=GtG.dtype).mul_(reg))
+            JJT = J @ J.T # (nresiduals, nresiduals)
+            if reg != 0:
+                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

@@ -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']