torchzero 0.1.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

torchzero/modules/second_order/newton.py

@@ -1,165 +1,142 @@
-from typing import Literal
-from collections import abc
-
-import torch
-
-from ...utils.derivatives import hessian_list_to_mat, jacobian_and_hessian
-from ...tensorlist import TensorList
-from ...core import OptimizerModule
-
-
-def _cholesky_solve(hessian: torch.Tensor, grad: torch.Tensor):
-    cholesky, info = torch.linalg.cholesky_ex(hessian) # pylint:disable=not-callable
-    if info == 0:
-        grad.unsqueeze_(1)
-        return torch.cholesky_solve(grad, cholesky), True
-    return None, False
-
-def _lu_solve(hessian: torch.Tensor, grad: torch.Tensor):
-    try:
-        newton_step, info = torch.linalg.solve_ex(hessian, grad) # pylint:disable=not-callable
-        if info == 0: return newton_step, True
-        return None, False
-    except torch.linalg.LinAlgError:
-        return None, False
-
-
-def _cholesky_fallback_lu(hessian: torch.Tensor, grad: torch.Tensor):
-    step, success = _cholesky_solve(hessian, grad)
-    if not success:
-        step, success = _lu_solve(hessian, grad)
-    return step, success
-
-def _least_squares_solve(hessian: torch.Tensor, grad: torch.Tensor):
-    return torch.linalg.lstsq(hessian, grad)[0], True # pylint:disable=not-callable
-
-
-def _fallback_gd(hessian:torch.Tensor, grad:torch.Tensor, lr = 1e-2):
-    return grad.mul_(1e-2), True
-
-def _fallback_safe_diag(hessian:torch.Tensor, grad:torch.Tensor, lr = 1e-2):
-    diag = hessian.diag().reciprocal_().nan_to_num_(1,1,1)
-    if torch.all(diag == 1): # fallback to gd
-        return _fallback_gd(hessian, grad, lr)
-    return grad.mul_(diag * lr), True
-
-
-def regularize_hessian_(hessian: torch.Tensor, value: float | Literal['eig']):
-    """regularize hessian matrix in-place"""
-    if value == 'eig':
-        value = torch.linalg.eigvalsh(hessian).min().clamp_(max=0).neg_() # pylint:disable=not-callable
-    elif value != 0:
-        hessian.add_(torch.eye(hessian.shape[0], device=hessian.device,dtype=hessian.dtype), alpha = value)
-
-LinearSystemSolvers = Literal['cholesky', 'lu', 'cholesky_lu', 'lstsq']
-FallbackLinearSystemSolvers = Literal['lstsq', 'safe_diag', 'gd']
-
-LINEAR_SYSTEM_SOLVERS = {
-    "cholesky": _cholesky_solve,
-    "lu": _lu_solve,
-    "cholesky_lu": _cholesky_fallback_lu,
-    "lstsq": _least_squares_solve,
-    "safe_diag": _fallback_safe_diag,
-    "gd": _fallback_gd
-}
-
-class ExactNewton(OptimizerModule):
-    """Peforms an exact Newton step using batched autograd.
-
-    Note that this doesn't support per-group settings.
-
-    Args:
-        tikhonov (float, optional):
-            tikhonov regularization (constant value added to the diagonal of the hessian).
-            Also known as Levenberg-Marquardt regularization. Can be set to 'eig', so it will be set
-            to the smallest eigenvalue of the hessian if that value is negative. Defaults to 0.
-        solver (Solvers, optional):
-            solver for Hx = g. Defaults to "cholesky_lu" (cholesky or LU if it fails).
-        fallback (Solvers, optional):
-            what to do if solver fails. Defaults to "safe_diag"
-            (takes nonzero diagonal elements, or fallbacks to gradient descent if all elements are 0).
-        validate (bool, optional):
-            validate if the step didn't increase the loss by `loss * tol` with an additional forward pass.
-            If not, undo the step and perform a gradient descent step.
-        tol (float, optional):
-            only has effect if `validate` is enabled.
-            If loss increased by `loss * tol`, perform gradient descent step.
-            Set this to 0 to guarantee that loss always decreases. Defaults to 1.
-        gd_lr (float, optional):
-            only has effect if `validate` is enabled.
-            Gradient descent step learning rate. Defaults to 1e-2.
-        batched_hessian (bool, optional):
-            whether to use experimental pytorch vmap-vectorized hessian calculation. As per pytorch docs,
-            should be faster, but this feature being experimental, there may be performance cliffs.
-            Defaults to True.
-        diag (False, optional):
-            only use the diagonal of the hessian. This will still calculate the full hessian!
-            This is mainly useful for benchmarking.
-    """
-    def __init__(
-        self,
-        tikhonov: float | Literal['eig'] = 0.0,
-        solver: LinearSystemSolvers = "cholesky_lu",
-        fallback: FallbackLinearSystemSolvers = "safe_diag",
-        validate=False,
-        tol: float = 1,
-        gd_lr = 1e-2,
-        batched_hessian=True,
-        diag: bool = False,
-    ):
-        super().__init__({})
-        self.tikhonov: float | Literal['eig'] = tikhonov
-        self.batched_hessian = batched_hessian
-
-        self.solver: abc.Callable = LINEAR_SYSTEM_SOLVERS[solver]
-        self.fallback: abc.Callable = LINEAR_SYSTEM_SOLVERS[fallback]
-
-        self.validate = validate
-        self.gd_lr = gd_lr
-        self.tol = tol
-
-        self.diag = diag
-
-    @torch.no_grad
-    def step(self, vars):
-        if vars.closure is None: raise ValueError("Newton requires a closure to compute the gradient.")
-
-        params = self.get_params()
-
-        # exact hessian via autograd
-        with torch.enable_grad():
-            vars.fx0 = vars.closure(False)
-            grads, hessian = jacobian_and_hessian([vars.fx0], params) # type:ignore
-            vars.grad = grads = TensorList(grads).squeeze_(0)
-            gvec = grads.to_vec()
-            hessian = hessian_list_to_mat(hessian)
-
-        # tikhonov regularization
-        regularize_hessian_(hessian, self.tikhonov)
-
-        # calculate newton step
-        if self.diag:
-            newton_step = gvec / hessian.diag()
-        else:
-            newton_step, success = self.solver(hessian, gvec)
-            if not success:
-                newton_step, success = self.fallback(hessian, gvec)
-                if not success:
-                    newton_step, success = _fallback_gd(hessian, gvec)
-
-        # apply the `_update` method
-        vars.ascent = grads.from_vec(newton_step.squeeze_().nan_to_num_(0,0,0))
-
-        # validate if newton step decreased loss
-        if self.validate:
-
-            params.sub_(vars.ascent)
-            fx1 = vars.closure(False)
-            params.add_(vars.ascent)
-
-            # if loss increases, set ascent direction to grad times lr
-            if (not fx1.isfinite()) or fx1 - vars.fx0 > vars.fx0 * self.tol: # type:ignore
-                vars.ascent = grads.div_(grads.total_vector_norm(2) / self.gd_lr)
-
-        # peform an update with the ascent direction, or pass it to the child.
-        return self._update_params_or_step_with_next(vars, params=params)
+import warnings
+from functools import partial
+from typing import Literal
+from collections.abc import Callable
+import torch
+
+from ...core import Chainable, apply, Module
+from ...utils import vec_to_tensors, TensorList
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    hessian_mat,
+    jacobian_and_hessian_wrt,
+)
+
+
+def lu_solve(H: torch.Tensor, g: torch.Tensor):
+    x, info = torch.linalg.solve_ex(H, g) # pylint:disable=not-callable
+    if info == 0: return x
+    return None
+
+def cholesky_solve(H: torch.Tensor, g: torch.Tensor):
+    x, info = torch.linalg.cholesky_ex(H) # pylint:disable=not-callable
+    if info == 0:
+        g.unsqueeze_(1)
+        return torch.cholesky_solve(g, x)
+    return None
+
+def least_squares_solve(H: torch.Tensor, g: torch.Tensor):
+    return torch.linalg.lstsq(H, g)[0] # pylint:disable=not-callable
+
+def eigh_solve(H: torch.Tensor, g: torch.Tensor, tfm: Callable | None):
+    try:
+        L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
+        if tfm is not None: L = tfm(L)
+        L.reciprocal_()
+        return torch.linalg.multi_dot([Q * L.unsqueeze(-2), Q.mH, g]) # pylint:disable=not-callable
+    except torch.linalg.LinAlgError:
+        return None
+
+def tikhonov_(H: torch.Tensor, reg: float):
+    if reg!=0: H.add_(torch.eye(H.size(-1), dtype=H.dtype, device=H.device).mul_(reg))
+    return H
+
+def eig_tikhonov_(H: torch.Tensor, reg: float):
+    v = torch.linalg.eigvalsh(H).min().clamp_(max=0).neg_() + reg # pylint:disable=not-callable
+    return tikhonov_(H, v)
+
+
+class Newton(Module):
+    """Exact newton via autograd.
+
+    Args:
+        reg (float, optional): tikhonov regularizer value. Defaults to 1e-6.
+        eig_reg (bool, optional): whether to use largest negative eigenvalue as regularizer. Defaults to False.
+        hessian_method (str):
+            how to calculate hessian. Defaults to "autograd".
+        vectorize (bool, optional):
+            whether to enable vectorized hessian. Defaults to True.
+        inner (Chainable | None, optional): inner modules. Defaults to None.
+        H_tfm (Callable | None, optional):
+            optional hessian transforms, takes in two arguments - `(hessian, gradient)`.
+
+            must return a tuple: `(hessian, is_inverted)` with transformed hessian and a boolean value
+            which must be True if transform inverted the hessian and False otherwise. Defaults to None.
+        eigval_tfm (Callable | None, optional):
+            optional eigenvalues transform, for example :code:`torch.abs` or :code:`lambda L: torch.clip(L, min=1e-8)`.
+            If this is specified, eigendecomposition will be used to solve Hx = g.
+
+    """
+    def __init__(
+        self,
+        reg: float = 1e-6,
+        eig_reg: bool = False,
+        hessian_method: Literal["autograd", "func", "autograd.functional"] = "autograd",
+        vectorize: bool = True,
+        inner: Chainable | None = None,
+        H_tfm: Callable[[torch.Tensor, torch.Tensor], tuple[torch.Tensor, bool]] | None = None,
+        eigval_tfm: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    ):
+        defaults = dict(reg=reg, eig_reg=eig_reg, abs=abs,hessian_method=hessian_method, vectorize=vectorize, H_tfm=H_tfm, eigval_tfm=eigval_tfm)
+        super().__init__(defaults)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        reg = settings['reg']
+        eig_reg = settings['eig_reg']
+        hessian_method = settings['hessian_method']
+        vectorize = settings['vectorize']
+        H_tfm = settings['H_tfm']
+        eigval_tfm = settings['eigval_tfm']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        if hessian_method == 'autograd':
+            with torch.enable_grad():
+                loss = vars.loss = vars.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
+                vars.grad = g_list
+                H = hessian_list_to_mat(H_list)
+
+        elif hessian_method in ('func', 'autograd.functional'):
+            strat = 'forward-mode' if vectorize else 'reverse-mode'
+            with torch.enable_grad():
+                g_list = vars.get_grad(retain_graph=True)
+                H: torch.Tensor = hessian_mat(partial(closure, backward=False), params,
+                                method=hessian_method, vectorize=vectorize, outer_jacobian_strategy=strat) # pyright:ignore[reportAssignmentType]
+
+        else:
+            raise ValueError(hessian_method)
+
+        # -------------------------------- inner step -------------------------------- #
+        if 'inner' in self.children:
+            g_list = apply(self.children['inner'], list(g_list), params=params, grads=list(g_list), vars=vars)
+        g = torch.cat([t.view(-1) for t in g_list])
+
+        # ------------------------------- regulazition ------------------------------- #
+        if eig_reg: H = eig_tikhonov_(H, reg)
+        else: H = tikhonov_(H, reg)
+
+        # ----------------------------------- solve ---------------------------------- #
+        update = None
+        if H_tfm is not None:
+            H, is_inv = H_tfm(H, g)
+            if is_inv: update = H
+
+        if eigval_tfm is not None:
+            update = eigh_solve(H, g, eigval_tfm)
+
+        if update is None: update = cholesky_solve(H, g)
+        if update is None: update = lu_solve(H, g)
+        if update is None: update = least_squares_solve(H, g)
+
+        vars.update = vec_to_tensors(update, params)
+        return vars

torchzero/modules/second_order/newton_cg.py

@@ -0,0 +1,84 @@
+from collections.abc import Callable
+from typing import Literal, overload
+import warnings
+import torch
+
+from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel
+from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
+
+from ...core import Chainable, apply, Module
+from ...utils.linalg.solve import cg
+
+class NewtonCG(Module):
+    def __init__(
+        self,
+        maxiter=None,
+        tol=1e-3,
+        reg: float = 1e-8,
+        hvp_method: Literal["forward", "central", "autograd"] = "forward",
+        h=1e-3,
+        warm_start=False,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(tol=tol, maxiter=maxiter, reg=reg, hvp_method=hvp_method, h=h, warm_start=warm_start)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+        closure = vars.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']
+        h = settings['h']
+        warm_start = settings['warm_start']
+
+        # ---------------------- Hessian vector product function --------------------- #
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+
+            def H_mm(x):
+                with torch.enable_grad():
+                    return TensorList(hvp(params, grad, x, retain_graph=True))
+
+        else:
+
+            with torch.enable_grad():
+                grad = vars.get_grad()
+
+            if hvp_method == 'forward':
+                def H_mm(x):
+                    return TensorList(hvp_fd_forward(closure, params, x, h=h, g_0=grad, normalize=True)[1])
+
+            elif hvp_method == 'central':
+                def H_mm(x):
+                    return TensorList(hvp_fd_central(closure, params, x, h=h, normalize=True)[1])
+
+            else:
+                raise ValueError(hvp_method)
+
+
+        # -------------------------------- inner step -------------------------------- #
+        b = grad
+        if 'inner' in self.children:
+            b = as_tensorlist(apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars))
+
+        # ---------------------------------- run cg ---------------------------------- #
+        x0 = None
+        if warm_start: x0 = self.get_state('prev_x', params=params, cls=TensorList) # initialized to 0 which is default anyway
+        x = cg(A_mm=H_mm, b=as_tensorlist(b), x0_=x0, tol=tol, maxiter=maxiter, reg=reg)
+        if warm_start:
+            assert x0 is not None
+            x0.set_(x)
+
+        vars.update = x
+        return vars
+
+

torchzero/modules/second_order/nystrom.py

@@ -0,0 +1,168 @@
+from collections.abc import Callable
+from typing import Literal, overload
+import warnings
+import torch
+
+from ...utils import TensorList, as_tensorlist, generic_zeros_like, generic_vector_norm, generic_numel, vec_to_tensors
+from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
+
+from ...core import Chainable, apply, Module
+from ...utils.linalg.solve import nystrom_sketch_and_solve, nystrom_pcg
+
+class NystromSketchAndSolve(Module):
+    def __init__(
+        self,
+        rank: int,
+        reg: float = 1e-3,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        h=1e-3,
+        inner: Chainable | None = None,
+        seed: int | None = None,
+    ):
+        defaults = dict(rank=rank, reg=reg, hvp_method=hvp_method, h=h, seed=seed)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        rank = settings['rank']
+        reg = settings['reg']
+        hvp_method = settings['hvp_method']
+        h = settings['h']
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+        # ---------------------- Hessian vector product function --------------------- #
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+
+            def H_mm(x):
+                with torch.enable_grad():
+                    Hvp = hvp(params, grad, params.from_vec(x), retain_graph=True)
+                    return torch.cat([t.ravel() for t in Hvp])
+
+        else:
+
+            with torch.enable_grad():
+                grad = vars.get_grad()
+
+            if hvp_method == 'forward':
+                def H_mm(x):
+                    Hvp = hvp_fd_forward(closure, params, params.from_vec(x), h=h, g_0=grad, normalize=True)[1]
+                    return torch.cat([t.ravel() for t in Hvp])
+
+            elif hvp_method == 'central':
+                def H_mm(x):
+                    Hvp = hvp_fd_central(closure, params, params.from_vec(x), h=h, normalize=True)[1]
+                    return torch.cat([t.ravel() for t in Hvp])
+
+            else:
+                raise ValueError(hvp_method)
+
+
+        # -------------------------------- inner step -------------------------------- #
+        b = grad
+        if 'inner' in self.children:
+            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+
+        # ------------------------------ sketch&n&solve ------------------------------ #
+        x = nystrom_sketch_and_solve(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), rank=rank, reg=reg, generator=generator)
+        vars.update = vec_to_tensors(x, reference=params)
+        return vars
+
+
+
+class NystromPCG(Module):
+    def __init__(
+        self,
+        sketch_size: int,
+        maxiter=None,
+        tol=1e-3,
+        reg: float = 1e-6,
+        hvp_method: Literal["forward", "central", "autograd"] = "autograd",
+        h=1e-3,
+        inner: Chainable | None = None,
+        seed: int | None = None,
+    ):
+        defaults = dict(sketch_size=sketch_size, reg=reg, maxiter=maxiter, tol=tol, hvp_method=hvp_method, h=h, seed=seed)
+        super().__init__(defaults,)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    @torch.no_grad
+    def step(self, vars):
+        params = TensorList(vars.params)
+
+        closure = vars.closure
+        if closure is None: raise RuntimeError('NewtonCG requires closure')
+
+        settings = self.settings[params[0]]
+        sketch_size = settings['sketch_size']
+        maxiter = settings['maxiter']
+        tol = settings['tol']
+        reg = settings['reg']
+        hvp_method = settings['hvp_method']
+        h = settings['h']
+
+
+        seed = settings['seed']
+        generator = None
+        if seed is not None:
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = torch.Generator(params[0].device).manual_seed(seed)
+            generator = self.global_state['generator']
+
+
+        # ---------------------- Hessian vector product function --------------------- #
+        if hvp_method == 'autograd':
+            grad = vars.get_grad(create_graph=True)
+
+            def H_mm(x):
+                with torch.enable_grad():
+                    Hvp = hvp(params, grad, params.from_vec(x), retain_graph=True)
+                    return torch.cat([t.ravel() for t in Hvp])
+
+        else:
+
+            with torch.enable_grad():
+                grad = vars.get_grad()
+
+            if hvp_method == 'forward':
+                def H_mm(x):
+                    Hvp = hvp_fd_forward(closure, params, params.from_vec(x), h=h, g_0=grad, normalize=True)[1]
+                    return torch.cat([t.ravel() for t in Hvp])
+
+            elif hvp_method == 'central':
+                def H_mm(x):
+                    Hvp = hvp_fd_central(closure, params, params.from_vec(x), h=h, normalize=True)[1]
+                    return torch.cat([t.ravel() for t in Hvp])
+
+            else:
+                raise ValueError(hvp_method)
+
+
+        # -------------------------------- inner step -------------------------------- #
+        b = grad
+        if 'inner' in self.children:
+            b = apply(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, vars=vars)
+
+        # ------------------------------ sketch&n&solve ------------------------------ #
+        x = nystrom_pcg(A_mm=H_mm, b=torch.cat([t.ravel() for t in b]), sketch_size=sketch_size, reg=reg, tol=tol, maxiter=maxiter, x0_=None, generator=generator)
+        vars.update = vec_to_tensors(x, reference=params)
+        return vars
+
+

torchzero/modules/smoothing/__init__.py

@@ -1,5 +1,2 @@
-r"""
-Gradient smoothing and orthogonalization methods.
-"""
-from .laplacian_smoothing import LaplacianSmoothing, gradient_laplacian_smoothing_
-from .gaussian_smoothing import GaussianHomotopy
+from .laplacian import LaplacianSmoothing
+from .gaussian import GaussianHomotopy