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

torchzero/modules/experimental/expanded_lbfgs.py

@@ -0,0 +1,141 @@
+from collections import deque
+from operator import itemgetter
+import torch
+
+from ...core import Transform, Chainable, Module, Var, apply_transform
+from ...utils import TensorList, as_tensorlist, NumberList
+from ...modules.quasi_newton.lbfgs import _adaptive_damping, lbfgs, _lerp_params_update_
+
+class ExpandedLBFGS(Module):
+    """L-BFGS but uses differences between more pairs than just consequtive. Window size controls how far away the pairs are allowed to be.
+    """
+    def __init__(
+        self,
+        history_size=10,
+        window_size:int=3,
+        tol: float | None = 1e-10,
+        damping: bool = False,
+        init_damping=0.9,
+        eigval_bounds=(0.5, 50),
+        params_beta: float | None = None,
+        grads_beta: float | None = None,
+        update_freq = 1,
+        z_beta: float | None = None,
+        tol_reset: bool = False,
+        inner: Chainable | None = None,
+    ):
+        defaults = dict(history_size=history_size, window_size=window_size, tol=tol, damping=damping, init_damping=init_damping, eigval_bounds=eigval_bounds, params_beta=params_beta, grads_beta=grads_beta, update_freq=update_freq, z_beta=z_beta, tol_reset=tol_reset)
+        super().__init__(defaults)
+
+        self.global_state['s_history'] = deque(maxlen=history_size)
+        self.global_state['y_history'] = deque(maxlen=history_size)
+        self.global_state['sy_history'] = deque(maxlen=history_size)
+        self.global_state['p_history'] = deque(maxlen=window_size)
+        self.global_state['g_history'] = deque(maxlen=window_size)
+
+        if inner is not None:
+            self.set_child('inner', inner)
+
+    def reset(self):
+        self.state.clear()
+        self.global_state['step'] = 0
+        self.global_state['s_history'].clear()
+        self.global_state['y_history'].clear()
+        self.global_state['sy_history'].clear()
+        self.global_state['p_history'].clear()
+        self.global_state['g_history'].clear()
+
+    @torch.no_grad
+    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
+
+        # history of s and k
+        s_history: deque[TensorList] = self.global_state['s_history']
+        y_history: deque[TensorList] = self.global_state['y_history']
+        sy_history: deque[torch.Tensor] = self.global_state['sy_history']
+        p_history: deque[TensorList] = self.global_state['p_history']
+        g_history: deque[TensorList] = self.global_state['g_history']
+
+        tol, damping, init_damping, eigval_bounds, update_freq, z_beta, tol_reset = itemgetter(
+            'tol', 'damping', 'init_damping', 'eigval_bounds', 'update_freq', 'z_beta', 'tol_reset')(self.settings[params[0]])
+        params_beta, grads_beta = self.get_settings(params, 'params_beta', 'grads_beta')
+
+        l_params, l_update = _lerp_params_update_(self, params, update, params_beta, grads_beta)
+        prev_l_params, prev_l_grad = self.get_state(params, 'prev_l_params', 'prev_l_grad', cls=TensorList)
+
+        # 1st step - there are no previous params and grads, lbfgs will do normalized GD step
+        if step == 0:
+            s = None; y = None; ys = None
+        else:
+            s = l_params - prev_l_params
+            y = l_update - prev_l_grad
+            ys = s.dot(y)
+
+            if damping:
+                s, y, ys = _adaptive_damping(s, y, ys, init_damping=init_damping, eigval_bounds=eigval_bounds)
+
+        prev_l_params.copy_(l_params)
+        prev_l_grad.copy_(l_update)
+
+        # update effective preconditioning state
+        if step % update_freq == 0:
+            if ys is not None and ys > 1e-10:
+                assert s is not None and y is not None
+                s_history.append(s)
+                y_history.append(y)
+                sy_history.append(ys)
+
+            if len(p_history) > 1:
+                for p_i, g_i in zip(list(p_history)[:-1], list(g_history)[:-1]):
+                    s_i = l_params - p_i
+                    y_i = l_update - g_i
+                    ys_i = s_i.dot(y_i)
+
+                    if ys_i > 1e-10:
+                        if damping:
+                            s_i, y_i, ys_i = _adaptive_damping(s_i, y_i, ys_i, init_damping=init_damping, eigval_bounds=eigval_bounds)
+
+                        s_history.append(s_i)
+                        y_history.append(y_i)
+                        sy_history.append(ys_i)
+
+            p_history.append(l_params.clone())
+            g_history.append(l_update.clone())
+
+
+        # step with inner module before applying preconditioner
+        if self.children:
+            update = TensorList(apply_transform(self.children['inner'], tensors=update, params=params, grads=var.grad, var=var))
+
+        # tolerance on gradient difference to avoid exploding after converging
+        if tol is not None:
+            if y is not None and y.abs().global_max() <= tol:
+                var.update = update # may have been updated by inner module, probably makes sense to use it here?
+                if tol_reset: self.reset()
+                return var
+
+        # lerp initial H^-1 @ q guess
+        z_ema = None
+        if z_beta is not None:
+            z_ema = self.get_state(var.params, 'z_ema', cls=TensorList)
+
+        # precondition
+        dir = lbfgs(
+            tensors_=as_tensorlist(update),
+            s_history=s_history,
+            y_history=y_history,
+            sy_history=sy_history,
+            y=y,
+            sy=ys,
+            z_beta = z_beta,
+            z_ema = z_ema,
+            step=step
+        )
+
+        var.update = dir
+
+        return var
+

torchzero/modules/{projections → experimental}/fft.py

@@ -2,12 +2,12 @@ import torch
 
 from ...core import Chainable
 from ...utils import vec_to_tensors
-from .projection import Projection
+from ..projections import ProjectionBase
 
 
-class FFTProjection(Projection):
+class FFTProjection(ProjectionBase):
     # norm description copied from pytorch docstring
-    """Project update into Fourrier space of real-valued inputs.
+    """Project update into Fourier space of real-valued inputs.
 
     Args:
         modules (Chainable): modules that will optimize the projected update.
@@ -45,8 +45,8 @@ class FFTProjection(Projection):
         super().__init__(modules, project_update=project_update, project_params=project_params, project_grad=project_grad, defaults=defaults)
 
     @torch.no_grad
-    def project(self, tensors, var, current):
-        settings = self.settings[var.params[0]]
+    def project(self, tensors, params, grads, loss, states, settings, current):
+        settings = settings[0]
         one_d = settings['one_d']
         norm = settings['norm']
 
@@ -60,14 +60,14 @@ class FFTProjection(Projection):
         return [torch.view_as_real(torch.fft.rfftn(t, norm=norm)) if t.numel() > 1 else t for t in tensors] # pylint:disable=not-callable
 
     @torch.no_grad
-    def unproject(self, tensors, var, current):
-        settings = self.settings[var.params[0]]
+    def unproject(self, projected_tensors, params, grads, loss, projected_states, projected_settings, current):
+        settings = projected_settings[0]
         one_d = settings['one_d']
         norm = settings['norm']
 
         if one_d:
-            vec = torch.view_as_complex(tensors[0])
+            vec = torch.view_as_complex(projected_tensors[0])
             unprojected_vec = torch.fft.irfft(vec, n=self.global_state['length'], norm=norm) # pylint:disable=not-callable
-            return vec_to_tensors(unprojected_vec, reference=var.params)
+            return vec_to_tensors(unprojected_vec, reference=params)
 
-        return [torch.fft.irfftn(torch.view_as_complex(t.contiguous()), s=p.shape, norm=norm) if t.numel() > 1 else t for t, p in zip(tensors, var.params)] # pylint:disable=not-callable
+        return [torch.fft.irfftn(torch.view_as_complex(t.contiguous()), s=p.shape, norm=norm) if t.numel() > 1 else t for t, p in zip(projected_tensors, params)] # pylint:disable=not-callable

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, Var, apply_transform, 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(
@@ -43,32 +43,31 @@ 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_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
+    # 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,

torchzero/modules/experimental/newtonnewton.py

@@ -22,8 +22,9 @@ from ..second_order.newton import (
 
 
 class NewtonNewton(Module):
-    """
-    Method that I thought of and then it worked.
+    """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
 
@@ -34,6 +35,9 @@ class NewtonNewton(Module):
     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,
@@ -83,6 +87,6 @@ class NewtonNewton(Module):
                     if x is None: x = least_squares_solve(H, xp)
                     xp = x.squeeze()
 
-        var.update = vec_to_tensors(xp, params)
+        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
+

torchzero/modules/experimental/reduce_outward_lr.py

@@ -4,19 +4,19 @@ from ...core import Target, Transform
 from ...utils import TensorList, unpack_states, unpack_dicts
 
 class ReduceOutwardLR(Transform):
-    """
-    When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
+    """When update sign matches weight sign, the learning rate for that weight is multiplied by `mul`.
 
     This means updates that move weights towards zero have higher learning rates.
 
-    A note on this is that it sounded good but its really bad in practice.
+    .. warning::
+        This sounded good but after testing turns out it sucks.
     """
     def __init__(self, mul = 0.5, use_grad=False, invert=False, target: Target = 'update'):
         defaults = dict(mul=mul, use_grad=use_grad, invert=invert)
         super().__init__(defaults, uses_grad=use_grad, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         params = TensorList(params)
         tensors = TensorList(tensors)