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

torchzero/modules/second_order/newton_cg.py

@@ -1,16 +1,16 @@
-import warnings
-import math
-from typing import Literal, cast
+
 from operator import itemgetter
+from typing import Literal, cast
+
 import torch
 
-from ...core import Chainable, Module, apply_transform
-from ...utils import TensorList, as_tensorlist, tofloat
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-from ...utils.linalg.solve import cg, minres, find_within_trust_radius
+from ...core import Chainable, Transform, HVPMethod
+from ...utils import TensorList, tofloat, unpack_dicts, unpack_states
+from ...linalg.solve import cg, find_within_trust_radius, minres
 from ..trust_region.trust_region import default_radius
 
-class NewtonCG(Module):
+
+class NewtonCG(Transform):
     """Newton's method with a matrix-free conjugate gradient or minimial-residual solver.
 
     Notes:
@@ -37,17 +37,14 @@ class NewtonCG(Module):
         hvp_method (str, optional):
             Determines how Hessian-vector products are evaluated.
 
-            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
-              This requires creating a graph for the gradient.
-            - ``"forward"``: Use a forward finite difference formula to
-              approximate the HVP. This requires one extra gradient evaluation.
-            - ``"central"``: Use a central finite difference formula for a
-              more accurate HVP approximation. This requires two extra
-              gradient evaluations.
-            Defaults to "autograd".
+            - ``"autograd"`` - uses autograd hessian-vector products. If multiple hessian-vector products are evaluated, uses a for-loop.
+            - ``"fd_forward"`` - uses gradient finite difference approximation with a less accurate forward formula which requires one extra gradient evaluation per hessian-vector product.
+            - ``"fd_central"`` - uses gradient finite difference approximation with a more accurate central formula which requires two gradient evaluations per hessian-vector product.
+
+            For NewtonCG ``"batched_autograd"`` is equivalent to ``"autograd"``. Defaults to ``"autograd"``.
         h (float, optional):
-            The step size for finite differences if :code:`hvp_method` is
-            ``"forward"`` or ``"central"``. Defaults to 1e-3.
+            The step size for finite difference if ``hvp_method`` is
+            ``"fd_forward"`` or ``"fd_central"``. Defaults to 1e-3.
         warm_start (bool, optional):
             If ``True``, the conjugate gradient solver is initialized with the
             solution from the previous optimization step. This can accelerate
@@ -82,100 +79,72 @@ class NewtonCG(Module):
         maxiter: int | None = None,
         tol: float = 1e-8,
         reg: float = 1e-8,
-        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
-        solver: Literal['cg', 'minres', 'minres_npc'] = 'cg',
+        hvp_method: HVPMethod = "autograd",
+        solver: Literal['cg', 'minres'] = 'cg',
+        npc_terminate: bool = False,
         h: float = 1e-3, # tuned 1e-4 or 1e-3
         miniter:int = 1,
         warm_start=False,
+        warm_beta:float=0,
         inner: Chainable | None = None,
     ):
         defaults = locals().copy()
         del defaults['self'], defaults['inner']
-        super().__init__(defaults,)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        super().__init__(defaults, inner=inner)
 
         self._num_hvps = 0
         self._num_hvps_last_step = 0
 
     @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        settings = self.settings[params[0]]
-        tol = settings['tol']
-        reg = settings['reg']
-        maxiter = settings['maxiter']
-        hvp_method = settings['hvp_method']
-        solver = settings['solver'].lower().strip()
-        h = settings['h']
-        warm_start = settings['warm_start']
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        hvp_method = fs['hvp_method']
+        h = fs['h']
 
-        self._num_hvps_last_step = 0
         # ---------------------- Hessian vector product function --------------------- #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-
-            def H_mm(x):
-                self._num_hvps_last_step += 1
-                with torch.enable_grad():
-                    return TensorList(hvp(params, grad, x, retain_graph=True))
-
-        else:
-
-            with torch.enable_grad():
-                grad = var.get_grad()
-
-            if hvp_method == 'forward':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
-
-            elif hvp_method == 'central':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
-
-            else:
-                raise ValueError(hvp_method)
+        _, H_mv = objective.list_Hvp_function(hvp_method=hvp_method, h=h, at_x0=True)
+        objective.temp = H_mv
 
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        self._num_hvps_last_step = 0
+        H_mv = objective.poptemp()
 
-        # -------------------------------- 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)
+        fs = settings[0]
+        tol = fs['tol']
+        reg = fs['reg']
+        maxiter = fs['maxiter']
+        solver = fs['solver'].lower().strip()
+        warm_start = fs['warm_start']
+        npc_terminate = fs["npc_terminate"]
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
-        if warm_start: x0 = self.get_state(params, 'prev_x', cls=TensorList) # initialized to 0 which is default anyway
+        if warm_start:
+            x0 = unpack_states(states, objective.params, 'prev_x', cls=TensorList)
+
+        b = TensorList(objective.get_updates())
 
         if solver == 'cg':
-            d, _ = cg(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, miniter=self.defaults["miniter"],reg=reg)
+            d, _ = cg(A_mv=H_mv, b=b, x0=x0, tol=tol, maxiter=maxiter,
+                      miniter=fs["miniter"], reg=reg, npc_terminate=npc_terminate)
 
         elif solver == 'minres':
-            d = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=False)
-
-        elif solver == 'minres_npc':
-            d = minres(A_mm=H_mm, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=True)
+            d = minres(A_mv=H_mv, b=b, x0=x0, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=npc_terminate)
 
         else:
             raise ValueError(f"Unknown solver {solver}")
 
         if warm_start:
             assert x0 is not None
-            x0.copy_(d)
-
-        var.update = d
+            x0.lerp_(d, weight = 1-fs["warm_beta"])
 
+        objective.updates = d
         self._num_hvps += self._num_hvps_last_step
-        return var
+        return objective
 
 
-class NewtonCGSteihaug(Module):
+class NewtonCGSteihaug(Transform):
     """Newton's method with trust region and a matrix-free Steihaug-Toint conjugate gradient solver.
 
     Notes:
@@ -219,7 +188,7 @@ class NewtonCGSteihaug(Module):
             whether to terminate CG/MINRES whenever negative curvature is detected. Defaults to False.
 
         hvp_method (str, optional):
-            either "forward" to use forward formula which requires one backward pass per Hvp, or "central" to use a more accurate central formula which requires two backward passes. "forward" is usually accurate enough. Defaults to "forward".
+            either ``"fd_forward"`` to use forward formula which requires one backward pass per hessian-vector product, or ``"fd_central"`` to use a more accurate central formula which requires two backward passes. ``"fd_forward"`` is usually accurate enough. Defaults to ``"fd_forward"``.
         h (float, optional): finite difference step size. Defaults to 1e-3.
 
         inner (Chainable | None, optional):
@@ -261,7 +230,7 @@ class NewtonCGSteihaug(Module):
         npc_terminate: bool = False,
 
         # hvp settings
-        hvp_method: Literal["forward", "central"] = "central",
+        hvp_method: Literal["fd_forward", "fd_central"] = "fd_central",
         h: float = 1e-3, # tuned 1e-4 or 1e-3
 
         # inner
@@ -269,72 +238,51 @@ class NewtonCGSteihaug(Module):
     ):
         defaults = locals().copy()
         del defaults['self'], defaults['inner']
-        super().__init__(defaults,)
-
-        if inner is not None:
-            self.set_child('inner', inner)
+        super().__init__(defaults, inner=inner)
 
         self._num_hvps = 0
         self._num_hvps_last_step = 0
 
-    @torch.no_grad
-    def step(self, var):
-        params = TensorList(var.params)
-        closure = var.closure
-        if closure is None: raise RuntimeError('NewtonCG requires closure')
-
-        tol = self.defaults['tol'] * self.global_state.get('tol_mul', 1)
-        solver = self.defaults['solver'].lower().strip()
-
-        (reg, maxiter, hvp_method, h, max_attempts, boundary_tol,
-         eta, nplus, nminus, rho_good, rho_bad, init, npc_terminate,
-         miniter, max_history, adapt_tol) = itemgetter(
-             "reg", "maxiter", "hvp_method", "h", "max_attempts", "boundary_tol",
-             "eta", "nplus", "nminus", "rho_good", "rho_bad", "init", "npc_terminate",
-             "miniter", "max_history", "adapt_tol",
-        )(self.defaults)
 
-        self._num_hvps_last_step = 0
+    @torch.no_grad
+    def update_states(self, objective, states, settings):
+        fs = settings[0]
+        hvp_method = fs['hvp_method']
+        h = fs['h']
 
         # ---------------------- Hessian vector product function --------------------- #
-        if hvp_method == 'autograd':
-            grad = var.get_grad(create_graph=True)
-
-            def H_mm(x):
-                self._num_hvps_last_step += 1
-                with torch.enable_grad():
-                    return TensorList(hvp(params, grad, x, retain_graph=True))
-
-        else:
-
-            with torch.enable_grad():
-                grad = var.get_grad()
+        _, H_mv = objective.list_Hvp_function(hvp_method=hvp_method, h=h, at_x0=True)
+        objective.temp = H_mv
 
-            if hvp_method == 'forward':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
+    @torch.no_grad
+    def apply_states(self, objective, states, settings):
+        self._num_hvps_last_step = 0
 
-            elif hvp_method == 'central':
-                def H_mm(x):
-                    self._num_hvps_last_step += 1
-                    return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
+        H_mv = objective.poptemp()
+        params = TensorList(objective.params)
+        fs = settings[0]
 
-            else:
-                raise ValueError(hvp_method)
+        tol = fs['tol'] * self.global_state.get('tol_mul', 1)
+        solver = fs['solver'].lower().strip()
 
+        reg=fs["reg"]
+        maxiter=fs["maxiter"]
+        max_attempts=fs["max_attempts"]
+        init=fs["init"]
+        npc_terminate=fs["npc_terminate"]
+        miniter=fs["miniter"]
+        max_history=fs["max_history"]
+        adapt_tol=fs["adapt_tol"]
 
-        # -------------------------------- 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)
 
         # ------------------------------- trust region ------------------------------- #
         success = False
         d = None
-        x0 = [p.clone() for p in params]
+        orig_params = [p.clone() for p in params]
+        b = TensorList(objective.get_updates())
         solution = None
+        closure = objective.closure
+        assert closure is not None
 
         while not success:
             max_attempts -= 1
@@ -343,7 +291,7 @@ class NewtonCGSteihaug(Module):
             trust_radius = self.global_state.get('trust_radius', init)
 
             # -------------- make sure trust radius isn't too small or large ------------- #
-            finfo = torch.finfo(x0[0].dtype)
+            finfo = torch.finfo(orig_params[0].dtype)
             if trust_radius < finfo.tiny * 2:
                 trust_radius = self.global_state['trust_radius'] = init
                 if adapt_tol:
@@ -360,7 +308,7 @@ class NewtonCGSteihaug(Module):
             if d is None:
                 if solver == 'cg':
                     d, solution = cg(
-                        A_mm=H_mm,
+                        A_mv=H_mv,
                         b=b,
                         tol=tol,
                         maxiter=maxiter,
@@ -372,40 +320,40 @@ class NewtonCGSteihaug(Module):
                     )
 
                 elif solver == 'minres':
-                    d = minres(A_mm=H_mm, b=b, trust_radius=trust_radius, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=npc_terminate)
+                    d = minres(A_mv=H_mv, b=b, trust_radius=trust_radius, tol=tol, maxiter=maxiter, reg=reg, npc_terminate=npc_terminate)
 
                 else:
                     raise ValueError(f"unknown solver {solver}")
 
             # ---------------------------- update trust radius --------------------------- #
             self.global_state["trust_radius"], success = default_radius(
-                params=params,
-                closure=closure,
-                f=tofloat(var.get_loss(False)),
-                g=b,
-                H=H_mm,
-                d=d,
-                trust_radius=trust_radius,
-                eta=eta,
-                nplus=nplus,
-                nminus=nminus,
-                rho_good=rho_good,
-                rho_bad=rho_bad,
-                boundary_tol=boundary_tol,
-
-                init=init, # init isn't used because check_overflow=False
-                state=self.global_state, # not used
-                settings=self.defaults, # not used
-                check_overflow=False, # this is checked manually to adapt tolerance
+                params = params,
+                closure = closure,
+                f = tofloat(objective.get_loss(False)),
+                g = b,
+                H = H_mv,
+                d = d,
+                trust_radius = trust_radius,
+                eta = fs["eta"],
+                nplus = fs["nplus"],
+                nminus = fs["nminus"],
+                rho_good = fs["rho_good"],
+                rho_bad = fs["rho_bad"],
+                boundary_tol = fs["boundary_tol"],
+
+                init = cast(int, None), # init isn't used because check_overflow=False
+                state = cast(dict, None), # not used
+                settings = cast(dict, None), # not used
+                check_overflow = False, # this is checked manually to adapt tolerance
             )
 
         # --------------------------- assign new direction --------------------------- #
         assert d is not None
         if success:
-            var.update = d
+            objective.updates = d
 
         else:
-            var.update = params.zeros_like()
+            objective.updates = params.zeros_like()
 
         self._num_hvps += self._num_hvps_last_step
-        return var
+        return objective