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

torchzero/modules/restarts/__init__.py

@@ -0,0 +1,7 @@
+from .restars import (
+    BirginMartinezRestart,
+    PowellRestart,
+    RestartEvery,
+    RestartOnStuck,
+    RestartStrategyBase,
+)

torchzero/modules/restarts/restars.py

@@ -0,0 +1,253 @@
+from abc import ABC, abstractmethod
+from functools import partial
+from typing import final, Literal, cast
+
+import torch
+
+from ...core import Chainable, Module, Var
+from ...utils import TensorList
+from ..termination import TerminationCriteriaBase
+
+def _reset_except_self(optimizer, var, self: Module):
+    for m in optimizer.unrolled_modules: m.reset()
+
+class RestartStrategyBase(Module, ABC):
+    """Base class for restart strategies.
+
+    On each ``update``/``step`` this checks reset condition and if it is satisfied,
+    resets the modules before updating or stepping.
+    """
+    def __init__(self, defaults: dict | None = None, modules: Chainable | None = None):
+        if defaults is None: defaults = {}
+        super().__init__(defaults)
+        if modules is not None:
+            self.set_child('modules', modules)
+
+    @abstractmethod
+    def should_reset(self, var: Var) -> bool:
+        """returns whether reset should occur"""
+
+    def _reset_on_condition(self, var):
+        modules = self.children.get('modules', None)
+
+        if self.should_reset(var):
+            if modules is None:
+                var.post_step_hooks.append(partial(_reset_except_self, self=self))
+            else:
+                modules.reset()
+
+        return modules
+
+    @final
+    def update(self, var):
+        modules = self._reset_on_condition(var)
+        if modules is not None:
+            modules.update(var)
+
+    @final
+    def apply(self, var):
+        # don't check here because it was check in `update`
+        modules = self.children.get('modules', None)
+        if modules is None: return var
+        return modules.apply(var.clone(clone_update=False))
+
+    @final
+    def step(self, var):
+        modules = self._reset_on_condition(var)
+        if modules is None: return var
+        return modules.step(var.clone(clone_update=False))
+
+
+
+class RestartOnStuck(RestartStrategyBase):
+    """Resets the state when update (difference in parameters) is zero for multiple steps in a row.
+
+    Args:
+        modules (Chainable | None):
+            modules to reset. If None, resets all modules.
+        tol (float, optional):
+            step is considered failed when maximum absolute parameter difference is smaller than this. Defaults to None (uses twice the smallest respresentable number)
+        n_tol (int, optional):
+            number of failed consequtive steps required to trigger a reset. Defaults to 10.
+
+    """
+    def __init__(self, modules: Chainable | None, tol: float | None = None, n_tol: int = 10):
+        defaults = dict(tol=tol, n_tol=n_tol)
+        super().__init__(defaults, modules)
+
+    @torch.no_grad
+    def should_reset(self, var):
+        step = self.global_state.get('step', 0)
+        self.global_state['step'] = step + 1
+
+        params = TensorList(var.params)
+        tol = self.defaults['tol']
+        if tol is None: tol = torch.finfo(params[0].dtype).tiny * 2
+        n_tol = self.defaults['n_tol']
+        n_bad = self.global_state.get('n_bad', 0)
+
+        # calculate difference in parameters
+        prev_params = self.get_state(params, 'prev_params', cls=TensorList)
+        update = params - prev_params
+        prev_params.copy_(params)
+
+        # if update is too small, it is considered bad, otherwise n_bad is reset to 0
+        if step > 0:
+            if update.abs().global_max() <= tol:
+                n_bad += 1
+
+            else:
+                n_bad = 0
+
+        self.global_state['n_bad'] = n_bad
+
+        # no progress, reset
+        if n_bad >= n_tol:
+            self.global_state.clear()
+            return True
+
+        return False
+
+
+class RestartEvery(RestartStrategyBase):
+    """Resets the state every n steps
+
+    Args:
+        modules (Chainable | None):
+            modules to reset. If None, resets all modules.
+        steps (int | Literal["ndim"]):
+            number of steps between resets. "ndim" to use number of parameters.
+    """
+    def __init__(self, modules: Chainable | None, steps: int | Literal['ndim']):
+        defaults = dict(steps=steps)
+        super().__init__(defaults, modules)
+
+    def should_reset(self, var):
+        step = self.global_state.get('step', 0) + 1
+        self.global_state['step'] = step
+
+        n = self.defaults['steps']
+        if isinstance(n, str): n = sum(p.numel() for p in var.params if p.requires_grad)
+
+        # reset every n steps
+        if step % n == 0:
+            self.global_state.clear()
+            return True
+
+        return False
+
+class RestartOnTerminationCriteria(RestartStrategyBase):
+    def __init__(self, modules: Chainable | None, criteria: "TerminationCriteriaBase"):
+        super().__init__(None, modules)
+        self.set_child('criteria', criteria)
+
+    def should_reset(self, var):
+        criteria = cast(TerminationCriteriaBase, self.children['criteria'])
+        return criteria.should_terminate(var)
+
+class PowellRestart(RestartStrategyBase):
+    """Powell's two restarting criterions for conjugate gradient methods.
+
+    The restart clears all states of ``modules``.
+
+    Args:
+        modules (Chainable | None):
+            modules to reset. If None, resets all modules.
+        cond1 (float | None, optional):
+            criterion that checks for nonconjugacy of the search directions.
+            Restart is performed whenevr g^Tg_{k+1} >= cond1*||g_{k+1}||^2.
+            The default condition value of 0.2 is suggested by Powell. Can be None to disable that criterion.
+        cond2 (float | None, optional):
+            criterion that checks if direction is not effectively downhill.
+            Restart is performed if -1.2||g||^2 < d^Tg < -0.8||g||^2.
+            Defaults to 0.2. Can be None to disable that criterion.
+
+    Reference:
+        Powell, Michael James David. "Restart procedures for the conjugate gradient method." Mathematical programming 12.1 (1977): 241-254.
+    """
+    def __init__(self, modules: Chainable | None, cond1:float | None = 0.2, cond2:float | None = 0.2):
+        defaults=dict(cond1=cond1, cond2=cond2)
+        super().__init__(defaults, modules)
+
+    def should_reset(self, var):
+        g = TensorList(var.get_grad())
+        cond1 = self.defaults['cond1']; cond2 = self.defaults['cond2']
+
+        # -------------------------------- initialize -------------------------------- #
+        if 'initialized' not in self.global_state:
+            self.global_state['initialized'] = 0
+            g_prev = self.get_state(var.params, 'g_prev', init=g)
+            return False
+
+        g_g = g.dot(g)
+
+        reset = False
+        # ------------------------------- 1st condition ------------------------------ #
+        if cond1 is not None:
+            g_prev = self.get_state(var.params, 'g_prev', must_exist=True, cls=TensorList)
+            g_g_prev = g_prev.dot(g)
+
+            if g_g_prev.abs() >= cond1 * g_g:
+                reset = True
+
+        # ------------------------------- 2nd condition ------------------------------ #
+        if (cond2 is not None) and (not reset):
+            d_g = TensorList(var.get_update()).dot(g)
+            if (-1-cond2) * g_g < d_g < (-1 + cond2) * g_g:
+                reset = True
+
+        # ------------------------------ clear on reset ------------------------------ #
+        if reset:
+            self.global_state.clear()
+            self.clear_state_keys('g_prev')
+            return True
+
+        return False
+
+# this requires direction from inner module,
+# so both this and inner have to be a single module
+class BirginMartinezRestart(Module):
+    """the restart criterion for conjugate gradient methods designed by Birgin and Martinez.
+
+    This criterion restarts when when the angle between dk+1 and −gk+1 is not acute enough.
+
+    The restart clears all states of ``module``.
+
+    Args:
+        module (Module):
+            module to restart, should be a conjugate gradient or possibly a quasi-newton method.
+        cond (float, optional):
+            Restart is performed whenevr d^Tg > -cond*||d||*||g||.
+            The default condition value of 1e-3 is suggested by Birgin and Martinez.
+
+    Reference:
+        Birgin, Ernesto G., and José Mario Martínez. "A spectral conjugate gradient method for unconstrained optimization." Applied Mathematics & Optimization 43.2 (2001): 117-128.
+    """
+    def __init__(self, module: Module, cond:float = 1e-3):
+        defaults=dict(cond=cond)
+        super().__init__(defaults)
+
+        self.set_child("module", module)
+
+    def update(self, var):
+        module = self.children['module']
+        module.update(var)
+
+    def apply(self, var):
+        module = self.children['module']
+        var = module.apply(var.clone(clone_update=False))
+
+        cond = self.defaults['cond']
+        g = TensorList(var.get_grad())
+        d = TensorList(var.get_update())
+        d_g = d.dot(g)
+        d_norm = d.global_vector_norm()
+        g_norm = g.global_vector_norm()
+
+        # d in our case is same direction as g so it has a minus sign
+        if -d_g > -cond * d_norm * g_norm:
+            module.reset()
+            var.update = g.clone()
+            return var
+
+        return var

torchzero/modules/second_order/__init__.py

@@ -1,3 +1,4 @@
 from .newton import Newton, InverseFreeNewton
-from .newton_cg import NewtonCG, TruncatedNewtonCG
+from .newton_cg import NewtonCG, NewtonCGSteihaug
 from .nystrom import NystromSketchAndSolve, NystromPCG
+from .multipoint import SixthOrder3P, SixthOrder5P, TwoPointNewton, SixthOrder3PM2

torchzero/modules/second_order/multipoint.py

@@ -0,0 +1,238 @@
+from collections.abc import Callable
+from contextlib import nullcontext
+from abc import ABC, abstractmethod
+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,
+)
+
+class HigherOrderMethodBase(Module, ABC):
+    def __init__(self, defaults: dict | None = None, vectorize: bool = True):
+        self._vectorize = vectorize
+        super().__init__(defaults)
+
+    @abstractmethod
+    def one_iteration(
+        self,
+        x: torch.Tensor,
+        evaluate: Callable[[torch.Tensor, int], tuple[torch.Tensor, ...]],
+        var: Var,
+    ) -> torch.Tensor:
+        """"""
+
+    @torch.no_grad
+    def step(self, var):
+        params = TensorList(var.params)
+        x0 = params.clone()
+        closure = var.closure
+        if closure is None: raise RuntimeError('MultipointNewton requires closure')
+        vectorize = self._vectorize
+
+        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)
+
+        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
+
+def _inv(A: torch.Tensor, lstsq:bool) -> torch.Tensor:
+    if lstsq: return torch.linalg.pinv(A) # pylint:disable=not-callable
+    A_inv, info = torch.linalg.inv_ex(A) # pylint:disable=not-callable
+    if info == 0: return A_inv
+    return torch.linalg.pinv(A) # pylint:disable=not-callable
+
+def _solve(A: torch.Tensor, b: torch.Tensor, lstsq: bool) -> torch.Tensor:
+    if lstsq: return torch.linalg.lstsq(A, b).solution # pylint:disable=not-callable
+    sol, info = torch.linalg.solve_ex(A, b) # pylint:disable=not-callable
+    if info == 0: return sol
+    return torch.linalg.lstsq(A, b).solution # pylint:disable=not-callable
+
+# 3f 2J 3 solves
+def sixth_order_3p(x:torch.Tensor, f, f_j, lstsq:bool=False):
+    f_x, J_x = f_j(x)
+
+    y = x - _solve(J_x, f_x, lstsq=lstsq)
+    f_y, J_y = f_j(y)
+
+    z = y - _solve(J_y, f_y, lstsq=lstsq)
+    f_z = f(z)
+
+    return y - _solve(J_y, f_y+f_z, lstsq=lstsq)
+
+class SixthOrder3P(HigherOrderMethodBase):
+    """Sixth-order iterative method.
+
+    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):
+        defaults=dict(lstsq=lstsq)
+        super().__init__(defaults=defaults, vectorize=vectorize)
+
+    def one_iteration(self, x, evaluate, var):
+        settings = self.defaults
+        lstsq = settings['lstsq']
+        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)
+        return x - x_star
+
+# I don't think it works (I tested root finding with this and it goes all over the place)
+# I double checked it multiple times
+# def sixth_order_im1(x:torch.Tensor, f, f_j, lstsq:bool=False):
+#     f_x, J_x = f_j(x)
+#     J_x_inv = _inv(J_x, lstsq=lstsq)
+
+#     y = x - J_x_inv @ f_x
+#     f_y, J_y = f_j(y)
+
+#     z = x - 2 * _solve(J_x + J_y, f_x, lstsq=lstsq)
+#     f_z = f(z)
+
+#     I = torch.eye(J_y.size(0), device=J_y.device, dtype=J_y.dtype)
+#     term1 = (7/2)*I
+#     term2 = 4 * (J_x_inv@J_y)
+#     term3 = (3/2) * (J_x_inv @ (J_y@J_y))
+
+#     return z - (term1 - term2 + term3) @ J_x_inv @ f_z
+
+# class SixthOrderIM1(HigherOrderMethodBase):
+#     """sixth-order iterative method https://www.mdpi.com/2504-3110/8/3/133
+
+#     """
+#     def __init__(self, lstsq: bool=False, vectorize: bool = True):
+#         defaults=dict(lstsq=lstsq)
+#         super().__init__(defaults=defaults, vectorize=vectorize)
+
+#     def iteration(self, x, evaluate, var):
+#         settings = self.defaults
+#         lstsq = settings['lstsq']
+#         def f(x): return evaluate(x, 1)[1]
+#         def f_j(x): return evaluate(x, 2)[1:]
+#         x_star = sixth_order_im1(x, f, f_j, lstsq)
+#         return x - x_star
+
+# 5f 5J 3 solves
+def sixth_order_5p(x:torch.Tensor, f_j, lstsq:bool=False):
+    f_x, J_x = f_j(x)
+    y = x - _solve(J_x, f_x, lstsq=lstsq)
+
+    f_y, J_y = f_j(y)
+    f_xy2, J_xy2 = f_j((x + y) / 2)
+
+    A = J_x + 2*J_xy2 + J_y
+
+    z = y - 4*_solve(A, f_y, lstsq=lstsq)
+    f_z, J_z = f_j(z)
+
+    f_xz2, J_xz2 = f_j((x + z) / 2)
+    B = J_x + 2*J_xz2 + J_z
+
+    return z - 4*_solve(B, f_z, lstsq=lstsq)
+
+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):
+        defaults=dict(lstsq=lstsq)
+        super().__init__(defaults=defaults, vectorize=vectorize)
+
+    def one_iteration(self, x, evaluate, var):
+        settings = self.defaults
+        lstsq = settings['lstsq']
+        def f_j(x): return evaluate(x, 2)[1:]
+        x_star = sixth_order_5p(x, f_j, lstsq)
+        return x - x_star
+
+# 2f 1J 2 solves
+def two_point_newton(x: torch.Tensor, f, f_j, lstsq:bool=False):
+    """third order convergence"""
+    f_x, J_x = f_j(x)
+    y = x - _solve(J_x, f_x, lstsq=lstsq)
+    f_y = f(y)
+    return x - _solve(J_x, f_x + f_y, lstsq=lstsq)
+
+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):
+        defaults=dict(lstsq=lstsq)
+        super().__init__(defaults=defaults, vectorize=vectorize)
+
+    def one_iteration(self, x, evaluate, var):
+        settings = self.defaults
+        lstsq = settings['lstsq']
+        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)
+        return x - x_star
+
+#3f 2J 1inv
+def sixth_order_3pm2(x:torch.Tensor, f, f_j, lstsq:bool=False):
+    f_x, J_x = f_j(x)
+    J_x_inv = _inv(J_x, lstsq=lstsq)
+    y = x - J_x_inv @ f_x
+    f_y, J_y = f_j(y)
+
+    I = torch.eye(x.numel(), dtype=x.dtype, device=x.device)
+    term = (2*I - J_x_inv @ J_y) @ J_x_inv
+    z = y - term @ f_y
+
+    return z - term @ f(z)
+
+
+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):
+        defaults=dict(lstsq=lstsq)
+        super().__init__(defaults=defaults, vectorize=vectorize)
+
+    def one_iteration(self, x, evaluate, var):
+        settings = self.defaults
+        lstsq = settings['lstsq']
+        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)
+        return x - x_star
+