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

torchzero/modules/second_order/multipoint.py

@@ -1,19 +1,17 @@
-from collections.abc import Callable
-from contextlib import nullcontext
 from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
+from typing import Any
+
 import numpy as np
 import torch
 
-from ...core import Chainable, Module, apply_transform, Var
-from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
-from ...utils.derivatives import (
-    flatten_jacobian,
-    jacobian_wrt,
-)
+from ...core import Chainable, DerivativesMethod, Objective, Transform
+from ...utils import TensorList, vec_to_tensors
+
 
-class HigherOrderMethodBase(Module, ABC):
-    def __init__(self, defaults: dict | None = None, vectorize: bool = True):
-        self._vectorize = vectorize
+class HigherOrderMethodBase(Transform, ABC):
+    def __init__(self, defaults: dict | None = None, derivatives_method: DerivativesMethod = 'batched_autograd'):
+        self._derivatives_method: DerivativesMethod = derivatives_method
         super().__init__(defaults)
 
     @abstractmethod
@@ -21,61 +19,27 @@ class HigherOrderMethodBase(Module, ABC):
         self,
         x: torch.Tensor,
         evaluate: Callable[[torch.Tensor, int], tuple[torch.Tensor, ...]],
-        var: Var,
+        objective: Objective,
+        setting: Mapping[str, Any],
     ) -> torch.Tensor:
         """"""
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        x0 = params.clone()
-        closure = var.closure
+    def apply_states(self, objective, states, settings):
+        params = TensorList(objective.params)
+
+        closure = objective.closure
         if closure is None: raise RuntimeError('MultipointNewton requires closure')
-        vectorize = self._vectorize
+        derivatives_method = self._derivatives_method
 
         def evaluate(x, order) -> tuple[torch.Tensor, ...]:
             """order=0 - returns (loss,), order=1 - returns (loss, grad), order=2 - returns (loss, grad, hessian), etc."""
-            params.from_vec_(x)
-
-            if order == 0:
-                loss = closure(False)
-                params.copy_(x0)
-                return (loss, )
-
-            if order == 1:
-                with torch.enable_grad():
-                    loss = closure()
-                grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                params.copy_(x0)
-                return loss, torch.cat([g.ravel() for g in grad])
-
-            with torch.enable_grad():
-                loss = var.loss = var.loss_approx = closure(False)
-
-                g_list = torch.autograd.grad(loss, params, create_graph=True)
-                var.grad = list(g_list)
-
-                g = torch.cat([t.ravel() for t in g_list])
-                n = g.numel()
-                ret = [loss, g]
-                T = g # current derivatives tensor
-
-                # get all derivative up to order
-                for o in range(2, order + 1):
-                    is_last = o == order
-                    T_list = jacobian_wrt([T], params, create_graph=not is_last, batched=vectorize)
-                    with torch.no_grad() if is_last else nullcontext():
-                        # the shape is (ndim, ) * order
-                        T = flatten_jacobian(T_list).view(n, n, *T.shape[1:])
-                        ret.append(T)
-
-            params.copy_(x0)
-            return tuple(ret)
+            return objective.derivatives_at(x, order, method=derivatives_method)
 
         x = torch.cat([p.ravel() for p in params])
-        dir = self.one_iteration(x, evaluate, var)
-        var.update = vec_to_tensors(dir, var.params)
-        return var
+        dir = self.one_iteration(x, evaluate, objective, settings[0])
+        objective.updates = vec_to_tensors(dir, objective.params)
+        return objective
 
 def _inv(A: torch.Tensor, lstsq:bool) -> torch.Tensor:
     if lstsq: return torch.linalg.pinv(A) # pylint:disable=not-callable
@@ -106,16 +70,15 @@ class SixthOrder3P(HigherOrderMethodBase):
 
     Abro, Hameer Akhtar, and Muhammad Mujtaba Shaikh. "A new time-efficient and convergent nonlinear solver." Applied Mathematics and Computation 355 (2019): 516-536.
     """
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f(x): return evaluate(x, 1)[1]
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = sixth_order_3p(x, f, f_j, lstsq)
+        x_star = sixth_order_3p(x, f, f_j, setting['lstsq'])
         return x - x_star
 
 # I don't think it works (I tested root finding with this and it goes all over the place)
@@ -173,15 +136,14 @@ def sixth_order_5p(x:torch.Tensor, f_j, lstsq:bool=False):
 
 class SixthOrder5P(HigherOrderMethodBase):
     """Argyros, Ioannis K., et al. "Extended convergence for two sixth order methods under the same weak conditions." Foundations 3.1 (2023): 127-139."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = sixth_order_5p(x, f_j, lstsq)
+        x_star = sixth_order_5p(x, f_j, setting['lstsq'])
         return x - x_star
 
 # 2f 1J 2 solves
@@ -196,16 +158,15 @@ class TwoPointNewton(HigherOrderMethodBase):
     """two-point Newton method with frozen derivative with third order convergence.
 
     Sharma, Janak Raj, and Deepak Kumar. "A fast and efficient composite Newton–Chebyshev method for systems of nonlinear equations." Journal of Complexity 49 (2018): 56-73."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f(x): return evaluate(x, 1)[1]
         def f_j(x): return evaluate(x, 2)[1:]
-        x_star = two_point_newton(x, f, f_j, lstsq)
+        x_star = two_point_newton(x, f, f_j, setting['lstsq'])
         return x - x_star
 
 #3f 2J 1inv
@@ -224,15 +185,14 @@ def sixth_order_3pm2(x:torch.Tensor, f, f_j, lstsq:bool=False):
 
 class SixthOrder3PM2(HigherOrderMethodBase):
     """Wang, Xiaofeng, and Yang Li. "An efficient sixth-order Newton-type method for solving nonlinear systems." Algorithms 10.2 (2017): 45."""
-    def __init__(self, lstsq: bool=False, vectorize: bool = True):
+    def __init__(self, lstsq: bool=False, derivatives_method: DerivativesMethod = 'batched_autograd'):
         defaults=dict(lstsq=lstsq)
-        super().__init__(defaults=defaults, vectorize=vectorize)
+        super().__init__(defaults=defaults, derivatives_method=derivatives_method)
 
-    def one_iteration(self, x, evaluate, var):
-        settings = self.defaults
-        lstsq = settings['lstsq']
+    @torch.no_grad
+    def one_iteration(self, x, evaluate, objective, setting):
         def f_j(x): return evaluate(x, 2)[1:]
         def f(x): return evaluate(x, 1)[1]
-        x_star = sixth_order_3pm2(x, f, f_j, lstsq)
+        x_star = sixth_order_3pm2(x, f, f_j, setting['lstsq'])
         return x - x_star
 

torchzero/modules/second_order/newton.py

@@ -1,21 +1,12 @@
-import warnings
 from collections.abc import Callable
-from functools import partial
 from typing import Literal
 
 import torch
 
-from ...core import Chainable, Module, apply_transform, Var
-from ...utils import TensorList, vec_to_tensors
-from ...utils.derivatives import (
-    flatten_jacobian,
-    hessian_mat,
-    hvp,
-    hvp_fd_central,
-    hvp_fd_forward,
-    jacobian_and_hessian_wrt,
-)
-from ...utils.linalg.linear_operator import DenseWithInverse, Dense
+from ...core import Chainable, Transform, Objective, HessianMethod, Module
+from ...utils import vec_to_tensors
+from ...linalg.linear_operator import Dense, DenseWithInverse
+
 
 def _lu_solve(H: torch.Tensor, g: torch.Tensor):
     try:
@@ -26,10 +17,9 @@ def _lu_solve(H: torch.Tensor, g: torch.Tensor):
         return None
 
 def _cholesky_solve(H: torch.Tensor, g: torch.Tensor):
-    x, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
+    L, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
     if info == 0:
-        g.unsqueeze_(1)
-        return torch.cholesky_solve(g, x)
+        return torch.cholesky_solve(g.unsqueeze(-1), L).squeeze(-1)
     return None
 
 def _least_squares_solve(H: torch.Tensor, g: torch.Tensor):
@@ -49,49 +39,14 @@ def _eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None, search_n
     except torch.linalg.LinAlgError:
         return None
 
-
-def _get_loss_grad_and_hessian(var: Var, hessian_method:str, vectorize:bool):
-    """returns (loss, g_list, H). Also sets var.loss and var.grad.
-    If hessian_method isn't 'autograd', loss is not set and returned as None"""
-    closure = var.closure
-    if closure is None:
-        raise RuntimeError("Second order methods requires a closure to be provided to the `step` method.")
-
-    params = var.params
-
-    # ------------------------ calculate grad and hessian ------------------------ #
-    loss = None
-    if hessian_method == 'autograd':
-        with torch.enable_grad():
-            loss = var.loss = var.loss_approx = closure(False)
-            g_list, H_list = jacobian_and_hessian_wrt([loss], params, batched=vectorize)
-            g_list = [t[0] for t in g_list] # remove leading dim from loss
-            var.grad = g_list
-            H = flatten_jacobian(H_list)
-
-    elif hessian_method in ('func', 'autograd.functional'):
-        strat = 'forward-mode' if vectorize else 'reverse-mode'
-        with torch.enable_grad():
-            g_list = var.get_grad(retain_graph=True)
-            H = hessian_mat(partial(closure, backward=False), params,
-                            method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
-
-    else:
-        raise ValueError(hessian_method)
-
-    return loss, g_list, H
-
-def _newton_step(var: Var, H: torch.Tensor, damping:float, inner: Module | None, H_tfm, eigval_fn, use_lstsq:bool, g_proj: Callable | None = None) -> torch.Tensor:
-    """returns the update tensor, then do vec_to_tensor(update, params)"""
-    params = var.params
-
-    if damping != 0:
-        H = H + torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping)
-
+def _newton_step(objective: Objective, H: torch.Tensor, damping:float, H_tfm, eigval_fn, use_lstsq:bool, g_proj: Callable | None = None, no_inner: Module | None = None) -> torch.Tensor:
+    """INNER SHOULD BE NONE IN MOST CASES! Because Transform already has inner.
+    Returns the update tensor, then do vec_to_tensor(update, params)"""
     # -------------------------------- inner step -------------------------------- #
-    update = var.get_update()
-    if inner is not None:
-        update = apply_transform(inner, update, params=params, grads=var.grad, loss=var.loss, var=var)
+    if no_inner is not None:
+        objective = no_inner.step(objective)
+
+    update = objective.get_updates()
 
     g = torch.cat([t.ravel() for t in update])
     if g_proj is not None: g = g_proj(g)
@@ -99,6 +54,9 @@ def _newton_step(var: Var, H: torch.Tensor, damping:float, inner: Module | None,
     # ----------------------------------- solve ---------------------------------- #
     update = None
 
+    if damping != 0:
+        H = H + torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(damping)
+
     if H_tfm is not None:
         ret = H_tfm(H, g)
 
@@ -133,7 +91,7 @@ def _get_H(H: torch.Tensor, eigval_fn):
 
     return Dense(H)
 
-class Newton(Module):
+class Newton(Transform):
     """Exact newton's method via autograd.
 
     Newton's method produces a direction jumping to the stationary point of quadratic approximation of the target function.
@@ -141,7 +99,7 @@ class Newton(Module):
     ``g`` can be output of another module, if it is specifed in ``inner`` argument.
 
     Note:
-        In most cases Newton should be the first module in the chain because it relies on autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
+        In most cases Newton should be the first module in the chain because it relies on autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
 
     Note:
         This module requires the a closure passed to the optimizer step,
@@ -158,10 +116,6 @@ class Newton(Module):
             when hessian is not invertible. If False, tries cholesky, if it fails tries LU, and then least squares.
             If ``eigval_fn`` is specified, eigendecomposition will always be used to solve the linear system and this
             argument will be ignored.
-        hessian_method (str):
-            how to calculate hessian. Defaults to "autograd".
-        vectorize (bool, optional):
-            whether to enable vectorized hessian. Defaults to True.
         H_tfm (Callable | None, optional):
             optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
 
@@ -174,6 +128,21 @@ class Newton(Module):
         eigval_fn (Callable | None, optional):
             optional eigenvalues transform, for example ``torch.abs`` or ``lambda L: torch.clip(L, min=1e-8)``.
             If this is specified, eigendecomposition will be used to invert the hessian.
+        hessian_method (str):
+            Determines how hessian is computed.
+
+            - ``"batched_autograd"`` - uses autograd to compute ``ndim`` batched hessian-vector products. Faster than ``"autograd"`` but uses more memory.
+            - ``"autograd"`` - uses autograd to compute ``ndim`` hessian-vector products using for loop. Slower than ``"batched_autograd"`` but uses less memory.
+            - ``"functional_revrev"`` - uses ``torch.autograd.functional`` with "reverse-over-reverse" strategy and a for-loop. This is generally equivalent to ``"autograd"``.
+            - ``"functional_fwdrev"`` - uses ``torch.autograd.functional`` with vectorized "forward-over-reverse" strategy. Faster than ``"functional_fwdrev"`` but uses more memory (``"batched_autograd"`` seems to be faster)
+            - ``"func"`` - uses ``torch.func.hessian`` which uses "forward-over-reverse" strategy. This method is the fastest and is recommended, however it is more restrictive and fails with some operators which is why it isn't the default.
+            - ``"gfd_forward"`` - computes ``ndim`` hessian-vector products via gradient finite difference using a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"gfd_central"`` - computes ``ndim`` hessian-vector products via gradient finite difference using a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+            - ``"fd"`` - uses function values to estimate gradient and hessian via finite difference. This uses less evaluations than chaining ``"gfd_*"`` after ``tz.m.FDM``.
+
+            Defaults to ``"batched_autograd"``.
+        h (float, optional):
+            finite difference step size for "fd_forward" and "fd_central".
         inner (Chainable | None, optional): modules to apply hessian preconditioner to. Defaults to None.
 
     # See also
@@ -249,45 +218,43 @@ class Newton(Module):
         damping: float = 0,
         use_lstsq: bool = False,
         update_freq: int = 1,
-        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
-        vectorize: bool = True,
         H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
         eigval_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+        hessian_method: HessianMethod = "batched_autograd",
+        h: float = 1e-3,
         inner: Chainable | None = None,
     ):
-        defaults = dict(damping=damping, hessian_method=hessian_method, use_lstsq=use_lstsq, vectorize=vectorize, H_tfm=H_tfm, eigval_fn=eigval_fn, update_freq=update_freq)
-        super().__init__(defaults)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        defaults = locals().copy()
+        del defaults['self'], defaults['update_freq'], defaults["inner"]
+        super().__init__(defaults, update_freq=update_freq, inner=inner)
 
     @torch.no_grad
-    def update(self, var):
-        step = self.global_state.get('step', 0)
-        self.global_state['step'] = step + 1
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
 
-        if step % self.defaults['update_freq'] == 0:
-            loss, g_list, self.global_state['H'] = _get_loss_grad_and_hessian(
-                var, self.defaults['hessian_method'], self.defaults['vectorize']
-            )
+        _, _, self.global_state['H'] = objective.hessian(
+            hessian_method=fs['hessian_method'],
+            h=fs['h'],
+            at_x0=True
+        )
 
     @torch.no_grad
-    def apply(self, var):
-        params = var.params
+    def apply_states(self, objective, states, settings):
+        params = objective.params
+        fs = settings[0]
+
         update = _newton_step(
-            var=var,
+            objective=objective,
             H = self.global_state["H"],
-            damping=self.defaults["damping"],
-            inner=self.children.get("inner", None),
-            H_tfm=self.defaults["H_tfm"],
-            eigval_fn=self.defaults["eigval_fn"],
-            use_lstsq=self.defaults["use_lstsq"],
+            damping = fs["damping"],
+            H_tfm = fs["H_tfm"],
+            eigval_fn = fs["eigval_fn"],
+            use_lstsq = fs["use_lstsq"],
         )
 
-        var.update = vec_to_tensors(update, params)
-
-        return var
+        objective.updates = vec_to_tensors(update, params)
+        return objective
 
-    def get_H(self,var=...):
+    def get_H(self,objective=...):
         return _get_H(self.global_state["H"], self.defaults["eigval_fn"])