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

torchzero/modules/experimental/hnewton.py

@@ -0,0 +1,85 @@
+from collections import deque
+
+import torch
+
+from ...core import TensorwiseTransform
+
+
+def eigh_solve(H: torch.Tensor, g: torch.Tensor):
+    try:
+        L, Q = torch.linalg.eigh(H) # pylint:disable=not-callable
+        return Q @ ((Q.mH @ g) / L)
+    except torch.linalg.LinAlgError:
+        return None
+
+
+class HNewton(TensorwiseTransform):
+    """This treats gradient differences as Hvps with vectors being parameter differences, using past gradients that are close to each other. Basically this is another limited memory quasi newton method to test.
+
+    .. warning::
+        Experimental.
+
+    """
+    def __init__(self, history_size: int, window_size: int, reg: float=0, tol: float = 1e-8, concat_params:bool=True, inner=None):
+        defaults = dict(history_size=history_size, window_size=window_size, reg=reg, tol=tol)
+        super().__init__(defaults, uses_grad=False, concat_params=concat_params, inner=inner)
+
+    def update_tensor(self, tensor, param, grad, loss, state, setting):
+
+        history_size = setting['history_size']
+
+        if 'param_history' not in state:
+            state['param_history'] = deque(maxlen=history_size)
+            state['grad_history'] = deque(maxlen=history_size)
+
+        param_history: deque = state['param_history']
+        grad_history: deque = state['grad_history']
+        param_history.append(param.ravel())
+        grad_history.append(tensor.ravel())
+
+    def apply_tensor(self, tensor, param, grad, loss, state, setting):
+        window_size = setting['window_size']
+        reg = setting['reg']
+        tol = setting['tol']
+
+        param_history: deque = state['param_history']
+        grad_history: deque = state['grad_history']
+        g = tensor.ravel()
+
+        n = len(param_history)
+        s_list = []
+        y_list = []
+
+        for i in range(n):
+            for j in range(i):
+                if i - j <= window_size:
+                    p_i, g_i = param_history[i], grad_history[i]
+                    p_j, g_j = param_history[j], grad_history[j]
+                    s = p_i - p_j # vec in hvp
+                    y = g_i - g_j # hvp
+                    if s.dot(y) > tol:
+                        s_list.append(s)
+                        y_list.append(y)
+
+        if len(s_list) < 1:
+            scale = (1 / tensor.abs().sum()).clip(min=torch.finfo(tensor.dtype).eps, max=1)
+            tensor.mul_(scale)
+            return tensor
+
+        S = torch.stack(s_list, 1)
+        Y = torch.stack(y_list, 1)
+
+        B = S.T @ Y
+        if reg != 0: B.add_(torch.eye(B.size(0), device=B.device, dtype=B.dtype).mul_(reg))
+        g_proj = g @ S
+
+        newton_proj, info = torch.linalg.solve_ex(B, g_proj) # pylint:disable=not-callable
+        if info != 0:
+            newton_proj = -torch.linalg.lstsq(B, g_proj).solution # pylint:disable=not-callable
+        newton = S @ newton_proj
+        return newton.view_as(tensor)
+
+
+        # scale = (1 / tensor.abs().sum()).clip(min=torch.finfo(tensor.dtype).eps, max=1)
+        # tensor.mul_(scale)
+        # return tensor

torchzero/modules/{quasi_newton/experimental → experimental}/modular_lbfgs.py

@@ -4,8 +4,8 @@ from typing import Any
 
 import torch
 
-from ....core import Chainable, Module, Transform, Vars, apply, maybe_chain
-from ....utils import NumberList, TensorList, as_tensorlist
+from ...core import Chainable, Module, Transform, Var, apply_transform, maybe_chain
+from ...utils import NumberList, TensorList, as_tensorlist
 
 
 def _adaptive_damping(
@@ -28,7 +28,7 @@ def _adaptive_damping(
 
 def lbfgs(
     tensors_: TensorList,
-    vars: Vars,
+    var: Var,
     s_history: deque[TensorList],
     y_history: deque[TensorList],
     sy_history: deque[torch.Tensor],
@@ -43,58 +43,57 @@ def lbfgs(
         if scale < 1e-5: scale = 1 / tensors_.abs().mean()
         return tensors_.mul_(min(1.0, scale)) # pyright: ignore[reportArgumentType]
 
-    else:
-        # 1st loop
-        alpha_list = []
-        q = tensors_.clone()
-        for s_i, y_i, ys_i in zip(reversed(s_history), reversed(y_history), reversed(sy_history)):
-            p_i = 1 / ys_i # this is also denoted as ρ (rho)
-            alpha = p_i * s_i.dot(q)
-            alpha_list.append(alpha)
-            q.sub_(y_i, alpha=alpha) # pyright: ignore[reportArgumentType]
-
-        # calculate z
-        # s.y/y.y is also this weird y-looking symbol I couldn't find
-        # z is it times q
-        # actually H0 = (s.y/y.y) * I, and z = H0 @ q
-        z = q * (ys_k / (y_k.dot(y_k)))
-
-        if z_tfm is not None:
-            z = TensorList(apply(z_tfm, tensors=z, params=vars.params, grads=vars.grad, vars=vars))
-
-        # 2nd loop
-        for s_i, y_i, ys_i, alpha_i in zip(s_history, y_history, sy_history, reversed(alpha_list)):
-            p_i = 1 / ys_i
-            beta_i = p_i * y_i.dot(z)
-            z.add_(s_i, alpha = alpha_i - beta_i)
-
-        return z
+    # 1st loop
+    alpha_list = []
+    q = tensors_.clone()
+    for s_i, y_i, ys_i in zip(reversed(s_history), reversed(y_history), reversed(sy_history)):
+        p_i = 1 / ys_i # this is also denoted as ρ (rho)
+        alpha = p_i * s_i.dot(q)
+        alpha_list.append(alpha)
+        q.sub_(y_i, alpha=alpha) # pyright: ignore[reportArgumentType]
+
+    # calculate z
+    # s.y/y.y is also this weird y-looking symbol I couldn't find
+    # z is it times q
+    # actually H0 = (s.y/y.y) * I, and z = H0 @ q
+    z = q * (ys_k / (y_k.dot(y_k)))
+
+    if z_tfm is not None:
+        z = TensorList(apply_transform(z_tfm, tensors=z, params=var.params, grads=var.grad, var=var))
+
+    # 2nd loop
+    for s_i, y_i, ys_i, alpha_i in zip(s_history, y_history, sy_history, reversed(alpha_list)):
+        p_i = 1 / ys_i
+        beta_i = p_i * y_i.dot(z)
+        z.add_(s_i, alpha = alpha_i - beta_i)
+
+    return z
 
 def _apply_tfms_into_history(
     self: Module,
     params: list[torch.Tensor],
-    vars: Vars,
+    var: Var,
     update: list[torch.Tensor],
 ):
     if 'params_history_tfm' in self.children:
-        params = apply(self.children['params_history_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=vars.grad, vars=vars)
+        params = apply_transform(self.children['params_history_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
 
     if 'grad_history_tfm' in self.children:
-        update = apply(self.children['grad_history_tfm'], tensors=as_tensorlist(update).clone(), params=params, grads=vars.grad, vars=vars)
+        update = apply_transform(self.children['grad_history_tfm'], tensors=as_tensorlist(update).clone(), params=params, grads=var.grad, var=var)
 
     return params, update
 
 def _apply_tfms_into_precond(
     self: Module,
     params: list[torch.Tensor],
-    vars: Vars,
+    var: Var,
     update: list[torch.Tensor],
 ):
     if 'params_precond_tfm' in self.children:
-        params = apply(self.children['params_precond_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=vars.grad, vars=vars)
+        params = apply_transform(self.children['params_precond_tfm'], tensors=as_tensorlist(params).clone(), params=params, grads=var.grad, var=var)
 
     if 'grad_precond_tfm' in self.children:
-        update = apply(self.children['grad_precond_tfm'], tensors=update, params=params, grads=vars.grad, vars=vars)
+        update = apply_transform(self.children['grad_precond_tfm'], tensors=update, params=params, grads=var.grad, var=var)
 
     return params, update
 
@@ -165,9 +164,9 @@ class ModularLBFGS(Module):
         self.global_state['sy_history'].clear()
 
     @torch.no_grad
-    def step(self, vars):
-        params = as_tensorlist(vars.params)
-        update = as_tensorlist(vars.get_update())
+    def step(self, var):
+        params = as_tensorlist(var.params)
+        update = as_tensorlist(var.get_update())
         step = self.global_state.get('step', 0)
         self.global_state['step'] = step + 1
 
@@ -186,11 +185,11 @@ class ModularLBFGS(Module):
         params_h, update_h = _apply_tfms_into_history(
             self,
             params=params,
-            vars=vars,
+            var=var,
             update=update,
         )
 
-        prev_params_h, prev_grad_h = self.get_state('prev_params_h', 'prev_grad_h', params=params, cls=TensorList)
+        prev_params_h, prev_grad_h = self.get_state(params, 'prev_params_h', 'prev_grad_h', cls=TensorList)
 
         # 1st step - there are no previous params and grads, `lbfgs` will do normalized SGD step
         if step == 0:
@@ -217,16 +216,16 @@ class ModularLBFGS(Module):
         # step with inner module before applying preconditioner
         if 'update_precond_tfm' in self.children:
             update_precond_tfm = self.children['update_precond_tfm']
-            inner_vars = update_precond_tfm.step(vars.clone(clone_update=True))
-            vars.update_attrs_from_clone_(inner_vars)
-            tensors = inner_vars.update
+            inner_var = update_precond_tfm.step(var.clone(clone_update=True))
+            var.update_attrs_from_clone_(inner_var)
+            tensors = inner_var.update
             assert tensors is not None
         else:
             tensors = update.clone()
 
         # transforms into preconditioner
-        params_p, update_p = _apply_tfms_into_precond(self, params=params, vars=vars, update=update)
-        prev_params_p, prev_grad_p = self.get_state('prev_params_p', 'prev_grad_p', params=params, cls=TensorList)
+        params_p, update_p = _apply_tfms_into_precond(self, params=params, var=var, update=update)
+        prev_params_p, prev_grad_p = self.get_state(params, 'prev_params_p', 'prev_grad_p', cls=TensorList)
 
         if step == 0:
             s_k_p = None; y_k_p = None; ys_k_p = None
@@ -245,13 +244,13 @@ class ModularLBFGS(Module):
         # tolerance on gradient difference to avoid exploding after converging
         if tol is not None:
             if y_k_p is not None and y_k_p.abs().global_max() <= tol:
-                vars.update = update # may have been updated by inner module, probably makes sense to use it here?
-                return vars
+                var.update = update # may have been updated by inner module, probably makes sense to use it here?
+                return var
 
         # precondition
         dir = lbfgs(
             tensors_=as_tensorlist(tensors),
-            vars=vars,
+            var=var,
             s_history=s_history,
             y_history=y_history,
             sy_history=sy_history,
@@ -260,7 +259,7 @@ class ModularLBFGS(Module):
             z_tfm=self.children.get('z_tfm', None),
         )
 
-        vars.update = dir
+        var.update = dir
 
-        return vars
+        return var
 

torchzero/modules/experimental/newton_solver.py

@@ -3,13 +3,13 @@ from typing import Any, Literal, overload
 
 import torch
 
-from ...core import Chainable, Module, apply, Modular
+from ...core import Chainable, Module, apply_transform, Modular
 from ...utils import TensorList, as_tensorlist
 from ...utils.derivatives import hvp
 from ..quasi_newton import LBFGS
 
 class NewtonSolver(Module):
-    """Matrix free newton via with any custom solver (usually it is better to just use NewtonCG or NystromPCG is even better)"""
+    """Matrix free newton via with any custom solver (this is for testing, use NewtonCG or NystromPCG)"""
     def __init__(
         self,
         solver: Callable[[list[torch.Tensor]], Any] = lambda p: Modular(p, LBFGS()),
@@ -26,9 +26,9 @@ class NewtonSolver(Module):
             self.set_child('inner', inner)
 
     @torch.no_grad
-    def step(self, vars):
-        params = TensorList(vars.params)
-        closure = vars.closure
+    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]]
@@ -39,7 +39,7 @@ class NewtonSolver(Module):
         warm_start = settings['warm_start']
 
         # ---------------------- Hessian vector product function --------------------- #
-        grad = vars.get_grad(create_graph=True)
+        grad = var.get_grad(create_graph=True)
 
         def H_mm(x):
             with torch.enable_grad():
@@ -50,11 +50,11 @@ class NewtonSolver(Module):
         # -------------------------------- inner step -------------------------------- #
         b = as_tensorlist(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))
+            b = as_tensorlist(apply_transform(self.children['inner'], [g.clone() for g in grad], params=params, grads=grad, var=var))
 
         # ---------------------------------- run cg ---------------------------------- #
         x0 = None
-        if warm_start: x0 = self.get_state('prev_x', params=params, cls=TensorList) # initialized to 0 which is default anyway
+        if warm_start: x0 = self.get_state(params, 'prev_x', cls=TensorList) # initialized to 0 which is default anyway
         if x0 is None: x = b.zeros_like().requires_grad_(True)
         else: x = x0.clone().requires_grad_(True)
 
@@ -76,13 +76,13 @@ class NewtonSolver(Module):
                 assert loss is not None
                 if min(loss, loss/initial_loss) < tol: break
 
-        print(f'{loss = }')
+        # print(f'{loss = }')
 
         if warm_start:
             assert x0 is not None
             x0.copy_(x)
 
-        vars.update = x.detach()
-        return vars
+        var.update = x.detach()
+        return var
 
 

torchzero/modules/experimental/newtonnewton.py

@@ -0,0 +1,92 @@
+import itertools
+import warnings
+from collections.abc import Callable
+from contextlib import nullcontext
+from functools import partial
+from typing import Literal
+
+import torch
+
+from ...core import Chainable, Module, apply_transform
+from ...utils import TensorList, vec_to_tensors
+from ...utils.derivatives import (
+    hessian_list_to_mat,
+    jacobian_wrt,
+)
+from ..second_order.newton import (
+    cholesky_solve,
+    eigh_solve,
+    least_squares_solve,
+    lu_solve,
+)
+
+
+class NewtonNewton(Module):
+    """Applies Newton-like preconditioning to Newton step.
+
+    This is a method that I thought of and then it worked. Here is how it works:
+
+    1. Calculate newton step by solving Hx=g
+
+    2. Calculate jacobian of x wrt parameters and call it H2
+
+    3. Solve H2 x2 = x for x2.
+
+    4. Optionally, repeat (if order is higher than 3.)
+
+    Memory is n^order. It tends to converge faster on convex functions, but can be unstable on non-convex. Orders higher than 3 are usually too unsable and have little benefit.
+
+    3rd order variant can minimize some convex functions with up to 100 variables in less time than Newton's method,
+    this is if pytorch can vectorize hessian computation efficiently.
+    """
+    def __init__(
+        self,
+        reg: float = 1e-6,
+        order: int = 3,
+        search_negative: bool = False,
+        vectorize: bool = True,
+        eigval_tfm: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    ):
+        defaults = dict(order=order, reg=reg, vectorize=vectorize, eigval_tfm=eigval_tfm, search_negative=search_negative)
+        super().__init__(defaults)
+
+    @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]]
+        reg = settings['reg']
+        vectorize = settings['vectorize']
+        order = settings['order']
+        search_negative = settings['search_negative']
+        eigval_tfm = settings['eigval_tfm']
+
+        # ------------------------ calculate grad and hessian ------------------------ #
+        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)
+
+            xp = torch.cat([t.ravel() for t in g_list])
+            I = torch.eye(xp.numel(), dtype=xp.dtype, device=xp.device)
+
+            for o in range(2, order + 1):
+                is_last = o == order
+                H_list = jacobian_wrt([xp], params, create_graph=not is_last, batched=vectorize)
+                with torch.no_grad() if is_last else nullcontext():
+                    H = hessian_list_to_mat(H_list)
+                    if reg != 0: H = H + I * reg
+
+                    x = None
+                    if search_negative or (is_last and eigval_tfm is not None):
+                        x = eigh_solve(H, xp, eigval_tfm, search_negative=search_negative)
+                    if x is None: x = cholesky_solve(H, xp)
+                    if x is None: x = lu_solve(H, xp)
+                    if x is None: x = least_squares_solve(H, xp)
+                    xp = x.squeeze()
+
+        var.update = vec_to_tensors(xp.nan_to_num_(0,0,0), params)
+        return var
+

torchzero/modules/experimental/parabolic_search.py

@@ -0,0 +1,220 @@
+import math
+from collections.abc import Mapping
+from operator import itemgetter
+
+import torch
+
+from ...core import Module
+from ...utils import TensorList
+
+
+
+def adaptive_tracking(
+    f,
+    f_0,
+    f_1,
+    t_0,
+    maxiter: int
+):
+
+    t = t_0
+    f_t = f(t)
+
+    # backtrack
+    if f_t > f_0:
+        if f_1 > f_0: t = min(0.5, t_0/2)
+        while f_t > f_0:
+            maxiter -= 1
+            if maxiter < 0: return 0, f_0
+            t = t/2
+            f_t = f(t) if t!=1 else f_1
+        return t, f_t
+
+    # forwardtrack
+    f_prev = f_t
+    t *= 2
+    f_t = f(t)
+    if f_prev < f_t: return t/2, f_prev
+    while f_prev >= f_t:
+        maxiter -= 1
+        if maxiter < 0: return t, f_t
+        f_prev = f_t
+        t *= 2
+        f_t = f(t)
+    return t/2, f_prev
+
+
+
+class ParabolaSearch(Module):
+    """"""
+    def __init__(
+        self,
+        step_size: float = 1e-2,
+        adaptive: bool=True,
+        normalize: bool=False,
+        # method: str | None = None,
+        maxiter: int | None = 10,
+        # bracket=None,
+        # bounds=None,
+        # tol: float | None = None,
+        # options=None,
+    ):
+        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
+        defaults = dict(step_size=step_size, adaptive=adaptive, normalize=normalize, maxiter=maxiter)
+        super().__init__(defaults)
+
+        import scipy.optimize
+        self.scopt = scipy.optimize
+
+
+    @torch.no_grad
+    def step(self, var):
+        x_0 = TensorList(var.params)
+        closure = var.closure
+        assert closure is not None
+        settings = self.settings[x_0[0]]
+        step_size = settings['step_size']
+        adaptive = settings['adaptive']
+        normalize = settings['normalize']
+        maxiter = settings['maxiter']
+        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
+
+        grad = TensorList(var.get_grad())
+        f_0 = var.get_loss(False)
+
+        scale = 1
+        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
+        if adaptive: scale = grad.abs().mean().clip(min=1e-8)
+
+        # make step
+        v_0 = grad * (step_size/scale)
+        x_0 -= v_0
+        with torch.enable_grad():
+            f_1 = closure()
+            grad = x_0.grad
+
+        x_0 += v_0
+        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
+        v_1 = grad * (step_size/scale)
+        a = v_1 - v_0
+
+        def parabolic_objective(t: float):
+            nonlocal x_0
+
+            step = v_0*t + 0.5*a*t**2
+            x_0 -= step
+            value = closure(False)
+            x_0 += step
+            return value.detach().cpu()
+
+        prev_t = self.global_state.get('prev_t', 2)
+        t, f = adaptive_tracking(parabolic_objective, f_0=f_0, f_1=f_1, t_0=prev_t, maxiter=maxiter)
+        self.global_state['prev_t'] = t
+
+        # method, bracket, bounds, tol, options, maxiter = itemgetter(
+        #     'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
+
+        # if maxiter is not None:
+        #     options = dict(options) if isinstance(options, Mapping) else {}
+        #     options['maxiter'] = maxiter
+
+        # res = self.scopt.minimize_scalar(parabolic_objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
+        # t = res.x
+
+        var.update = v_0*t + 0.5*a*t**2
+        return var
+
+class CubicParabolaSearch(Module):
+    """"""
+    def __init__(
+        self,
+        step_size: float = 1e-2,
+        adaptive: bool=True,
+        normalize: bool=False,
+        # method: str | None = None,
+        maxiter: int | None = 10,
+        # bracket=None,
+        # bounds=None,
+        # tol: float | None = None,
+        # options=None,
+    ):
+        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
+        defaults = dict(step_size=step_size, adaptive=adaptive, normalize=normalize, maxiter=maxiter)
+        super().__init__(defaults)
+
+        import scipy.optimize
+        self.scopt = scipy.optimize
+
+
+    @torch.no_grad
+    def step(self, var):
+        x_0 = TensorList(var.params)
+        closure = var.closure
+        assert closure is not None
+        settings = self.settings[x_0[0]]
+        step_size = settings['step_size']
+        adaptive = settings['adaptive']
+        maxiter = settings['maxiter']
+        normalize = settings['normalize']
+        if normalize and adaptive: raise ValueError("pick either normalize or adaptive")
+
+        grad = TensorList(var.get_grad())
+        f_0 = var.get_loss(False)
+
+        scale = 1
+        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
+        if adaptive: scale = grad.abs().mean().clip(min=1e-8)
+
+        # make step
+        v_0 = grad * (step_size/scale)
+        x_0 -= v_0
+        with torch.enable_grad():
+            f_1 = closure()
+            grad = x_0.grad
+
+        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
+        v_1 = grad * (step_size/scale)
+        a_0 = v_1 - v_0
+
+        # make another step
+        x_0 -= v_1
+        with torch.enable_grad():
+            f_2 = closure()
+            grad = x_0.grad
+
+        if normalize: grad = grad/grad.abs().mean().clip(min=1e-8)
+        v_2 = grad * (step_size/scale)
+        a_1 = v_2 - v_1
+
+        j = a_1 - a_0
+
+        x_0 += v_0
+        x_0 += v_1
+
+        def parabolic_objective(t: float):
+            nonlocal x_0
+
+            step = v_0*t + (1/2)*a_0*t**2 + (1/6)*j*t**3
+            x_0 -= step
+            value = closure(False)
+            x_0 += step
+            return value
+
+
+        prev_t = self.global_state.get('prev_t', 2)
+        t, f = adaptive_tracking(parabolic_objective, f_0=f_0, f_1=f_1, t_0=prev_t, maxiter=maxiter)
+        self.global_state['prev_t'] = t
+
+        # method, bracket, bounds, tol, options, maxiter = itemgetter(
+        #     'method', 'bracket', 'bounds', 'tol', 'options', 'maxiter')(self.settings[var.params[0]])
+
+        # if maxiter is not None:
+        #     options = dict(options) if isinstance(options, Mapping) else {}
+        #     options['maxiter'] = maxiter
+
+        # res = self.scopt.minimize_scalar(parabolic_objective, method=method, bracket=bracket, bounds=bounds, tol=tol, options=options)
+        # t = res.x
+
+        var.update = v_0*t + (1/2)*a_0*t**2 + (1/6)*j*t**3
+        return var
+