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

torchzero/modules/higher_order/higher_order_newton.py

@@ -13,7 +13,7 @@ import torch
 from ...core import Chainable, Module, apply_transform
 from ...utils import TensorList, vec_to_tensors, vec_to_tensors_
 from ...utils.derivatives import (
-    hessian_list_to_mat,
+    flatten_jacobian,
     jacobian_wrt,
 )
 
@@ -148,21 +148,16 @@ class HigherOrderNewton(Module):
     """A basic arbitrary order newton's method with optional trust region and proximal penalty.
 
     This constructs an nth order taylor approximation via autograd and minimizes it with
-    scipy.optimize.minimize trust region newton solvers with optional proximal penalty.
+    ``scipy.optimize.minimize`` trust region newton solvers with optional proximal penalty.
 
-    .. note::
-        In most cases HigherOrderNewton should be the first module in the chain because it relies on extra autograd. Use the :code:`inner` argument if you wish to apply Newton preconditioning to another module's output.
+    The hessian of taylor approximation is easier to evaluate, plus it can be evaluated in a batched mode,
+    so it can be more efficient in very specific instances.
 
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating higher order derivatives.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    .. warning::
-        this uses roughly O(N^order) memory and solving the subproblem can be very expensive.
-
-    .. warning::
-        "none" and "proximal" trust methods may generate subproblems that have no minima, causing divergence.
+    Notes:
+        - In most cases HigherOrderNewton should be the first module in the chain because it relies on extra autograd. Use the ``inner`` argument if you wish to apply Newton preconditioning to another module's output.
+        - This module requires the a closure passed to the optimizer step, as it needs to re-evaluate the loss and gradients for calculating higher order derivatives. The closure must accept a ``backward`` argument (refer to documentation).
+        - this uses roughly O(N^order) memory and solving the subproblem is very expensive.
+        - "none" and "proximal" trust methods may generate subproblems that have no minima, causing divergence.
 
     Args:
 
@@ -178,7 +173,7 @@ class HigherOrderNewton(Module):
         increase (float, optional): trust region multiplier on good steps. Defaults to 1.5.
         decrease (float, optional): trust region multiplier on bad steps. Defaults to 0.75.
         trust_init (float | None, optional):
-            initial trust region size. If none, defaults to 1 on :code:`trust_method="bounds"` and 0.1 on :code:`"proximal"`. Defaults to None.
+            initial trust region size. If none, defaults to 1 on :code:`trust_method="bounds"` and 0.1 on ``"proximal"``. Defaults to None.
         trust_tol (float, optional):
             Maximum ratio of expected loss reduction to actual reduction for trust region increase.
             Should 1 or higer. Defaults to 2.
@@ -191,11 +186,14 @@ class HigherOrderNewton(Module):
         self,
         order: int = 4,
         trust_method: Literal['bounds', 'proximal', 'none'] | None = 'bounds',
-        nplus: float = 2,
+        nplus: float = 3.5,
         nminus: float = 0.25,
+        rho_good: float = 0.99,
+        rho_bad: float = 1e-4,
         init: float | None = None,
         eta: float = 1e-6,
         max_attempts = 10,
+        boundary_tol: float = 1e-2,
         de_iters: int | None = None,
         vectorize: bool = True,
     ):
@@ -203,7 +201,7 @@ class HigherOrderNewton(Module):
             if trust_method == 'bounds': init = 1
             else: init = 0.1
 
-        defaults = dict(order=order, trust_method=trust_method, nplus=nplus, nminus=nminus, eta=eta, init=init, vectorize=vectorize, de_iters=de_iters, max_attempts=max_attempts)
+        defaults = dict(order=order, trust_method=trust_method, nplus=nplus, nminus=nminus, eta=eta, init=init, vectorize=vectorize, de_iters=de_iters, max_attempts=max_attempts, boundary_tol=boundary_tol, rho_good=rho_good, rho_bad=rho_bad)
         super().__init__(defaults)
 
     @torch.no_grad
@@ -222,6 +220,9 @@ class HigherOrderNewton(Module):
         de_iters = settings['de_iters']
         max_attempts = settings['max_attempts']
         vectorize = settings['vectorize']
+        boundary_tol = settings['boundary_tol']
+        rho_good = settings['rho_good']
+        rho_bad = settings['rho_bad']
 
         # ------------------------ calculate grad and hessian ------------------------ #
         with torch.enable_grad():
@@ -241,7 +242,7 @@ class HigherOrderNewton(Module):
                 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 = hessian_list_to_mat(T_list).view(n, n, *T.shape[1:])
+                    T = flatten_jacobian(T_list).view(n, n, *T.shape[1:])
                     derivatives.append(T)
 
         x0 = torch.cat([p.ravel() for p in params])
@@ -254,8 +255,13 @@ class HigherOrderNewton(Module):
 
             # load trust region value
             trust_value = self.global_state.get('trust_region', init)
-            if trust_value < 1e-8 or trust_value > 1e16: trust_value = self.global_state['trust_region'] = settings['init']
 
+            # make sure its not too small or too large
+            finfo = torch.finfo(x0.dtype)
+            if trust_value < finfo.tiny*2 or trust_value > finfo.max / (2*nplus):
+                trust_value = self.global_state['trust_region'] = settings['init']
+
+            # determine tr and prox values
             if trust_method is None: trust_method = 'none'
             else: trust_method = trust_method.lower()
 
@@ -297,13 +303,15 @@ class HigherOrderNewton(Module):
 
                 rho = reduction / (max(pred_reduction, 1e-8))
                 # failed step
-                if rho < 0.25:
+                if rho < rho_bad:
                     self.global_state['trust_region'] = trust_value * nminus
 
                 # very good step
-                elif rho > 0.75:
-                    diff = trust_value - (x0 - x_star).abs_()
-                    if (diff.amin() / trust_value) > 1e-4: # hits boundary
+                elif rho > rho_good:
+                    step = (x_star - x0)
+                    magn = torch.linalg.vector_norm(step) # pylint:disable=not-callable
+                    if trust_method == 'proximal' or (trust_value - magn) / trust_value <= boundary_tol:
+                        # close to boundary
                         self.global_state['trust_region'] = trust_value * nplus
 
                 # if the ratio is high enough then accept the proposed step

torchzero/modules/least_squares/__init__.py

@@ -0,0 +1 @@
+from .gn import SumOfSquares, GaussNewton

torchzero/modules/least_squares/gn.py

@@ -0,0 +1,161 @@
+import torch
+from ...core import Module
+
+from ...utils.derivatives import jacobian_wrt, flatten_jacobian
+from ...utils import vec_to_tensors
+from ...utils.linalg import linear_operator
+class SumOfSquares(Module):
+    """Sets loss to be the sum of squares of values returned by the closure.
+
+    This is meant to be used to test least squares methods against ordinary minimization methods.
+
+    To use this, the closure should return a vector of values to minimize sum of squares of.
+    Please add the `backward` argument, it will always be False but it is required.
+    """
+    def __init__(self):
+        super().__init__()
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+
+        if closure is not None:
+            def sos_closure(backward=True):
+                if backward:
+                    var.zero_grad()
+                    with torch.enable_grad():
+                        loss = closure(False)
+                        loss = loss.pow(2).sum()
+                        loss.backward()
+                    return loss
+
+                loss = closure(False)
+                return loss.pow(2).sum()
+
+            var.closure = sos_closure
+
+        if var.loss is not None:
+            var.loss = var.loss.pow(2).sum()
+
+        if var.loss_approx is not None:
+            var.loss_approx = var.loss_approx.pow(2).sum()
+
+        return var
+
+
+class GaussNewton(Module):
+    """Gauss-newton method.
+
+    To use this, the closure should return a vector of values to minimize sum of squares of.
+    Please add the ``backward`` argument, it will always be False but it is required.
+    Gradients will be calculated via batched autograd within this module, you don't need to
+    implement the backward pass. Please see below for an example.
+
+    Note:
+        This method requires ``ndim^2`` memory, however, if it is used within ``tz.m.TrustCG`` trust region,
+        the memory requirement is ``ndim*m``, where ``m`` is number of values in the output.
+
+    Args:
+        reg (float, optional): regularization parameter. Defaults to 1e-8.
+        batched (bool, optional): whether to use vmapping. Defaults to True.
+
+    Examples:
+
+    minimizing the rosenbrock function:
+    ```python
+    def rosenbrock(X):
+        x1, x2 = X
+        return torch.stack([(1 - x1), 100 * (x2 - x1**2)])
+
+    X = torch.tensor([-1.1, 2.5], requires_grad=True)
+    opt = tz.Modular([X], tz.m.GaussNewton(), tz.m.Backtracking())
+
+    # define the closure for line search
+    def closure(backward=True):
+        return rosenbrock(X)
+
+    # minimize
+    for iter in range(10):
+        loss = opt.step(closure)
+        print(f'{loss = }')
+    ```
+
+    training a neural network with a matrix-free GN trust region:
+    ```python
+    X = torch.randn(64, 20)
+    y = torch.randn(64, 10)
+
+    model = nn.Sequential(nn.Linear(20, 64), nn.ELU(), nn.Linear(64, 10))
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.TrustCG(tz.m.GaussNewton()),
+    )
+
+    def closure(backward=True):
+        y_hat = model(X) # (64, 10)
+        return (y_hat - y).pow(2).mean(0) # (10, )
+
+    for i in range(100):
+        losses = opt.step(closure)
+        if i % 10 == 0:
+            print(f'{losses.mean() = }')
+    ```
+    """
+    def __init__(self, reg:float = 1e-8, batched:bool=True, ):
+        super().__init__(defaults=dict(batched=batched, reg=reg))
+
+    @torch.no_grad
+    def update(self, var):
+        params = var.params
+        batched = self.defaults['batched']
+
+        closure = var.closure
+        assert closure is not None
+
+        # gauss newton direction
+        with torch.enable_grad():
+            f = var.get_loss(backward=False) # n_out
+            assert isinstance(f, torch.Tensor)
+            G_list = jacobian_wrt([f.ravel()], params, batched=batched)
+
+        var.loss = f.pow(2).sum()
+
+        G = self.global_state["G"] = flatten_jacobian(G_list) # (n_out, ndim)
+        Gtf = G.T @ f.detach() # (ndim)
+        self.global_state["Gtf"] = Gtf
+        var.grad = vec_to_tensors(Gtf, var.params)
+
+        # set closure to calculate sum of squares for line searches etc
+        if var.closure is not None:
+            def sos_closure(backward=True):
+                if backward:
+                    var.zero_grad()
+                    with torch.enable_grad():
+                        loss = closure(False).pow(2).sum()
+                        loss.backward()
+                    return loss
+
+                loss = closure(False).pow(2).sum()
+                return loss
+
+            var.closure = sos_closure
+
+    @torch.no_grad
+    def apply(self, var):
+        reg = self.defaults['reg']
+
+        G = self.global_state['G']
+        Gtf = self.global_state['Gtf']
+
+        GtG = G.T @ G # (ndim, ndim)
+        if reg != 0:
+            GtG.add_(torch.eye(GtG.size(0), device=GtG.device, dtype=GtG.dtype).mul_(reg))
+
+        v = torch.linalg.lstsq(GtG, Gtf).solution # pylint:disable=not-callable
+
+        var.update = vec_to_tensors(v, var.params)
+        return var
+
+    def get_H(self, var):
+        G = self.global_state['G']
+        return linear_operator.AtA(G)

torchzero/modules/line_search/__init__.py

@@ -1,5 +1,5 @@
-from .adaptive import AdaptiveLineSearch
+from .adaptive import AdaptiveTracking
 from .backtracking import AdaptiveBacktracking, Backtracking
 from .line_search import LineSearchBase
 from .scipy import ScipyMinimizeScalar
-from .strong_wolfe import StrongWolfe
+from .strong_wolfe import StrongWolfe

torchzero/modules/line_search/_polyinterp.py

@@ -0,0 +1,289 @@
+import numpy as np
+import torch
+
+from .line_search import LineSearchBase
+
+
+# polynomial interpolation
+# this code is from https://github.com/hjmshi/PyTorch-LBFGS/blob/master/functions/LBFGS.py
+# PyTorch-LBFGS: A PyTorch Implementation of L-BFGS
+def polyinterp(points, x_min_bound=None, x_max_bound=None, plot=False):
+    """
+    Gives the minimizer and minimum of the interpolating polynomial over given points
+    based on function and derivative information. Defaults to bisection if no critical
+    points are valid.
+
+    Based on polyinterp.m Matlab function in minFunc by Mark Schmidt with some slight
+    modifications.
+
+    Implemented by: Hao-Jun Michael Shi and Dheevatsa Mudigere
+    Last edited 12/6/18.
+
+    Inputs:
+        points (nparray): two-dimensional array with each point of form [x f g]
+        x_min_bound (float): minimum value that brackets minimum (default: minimum of points)
+        x_max_bound (float): maximum value that brackets minimum (default: maximum of points)
+        plot (bool): plot interpolating polynomial
+
+    Outputs:
+        x_sol (float): minimizer of interpolating polynomial
+        F_min (float): minimum of interpolating polynomial
+
+    Note:
+      . Set f or g to np.nan if they are unknown
+
+    """
+    no_points = points.shape[0]
+    order = np.sum(1 - np.isnan(points[:, 1:3]).astype('int')) - 1
+
+    x_min = np.min(points[:, 0])
+    x_max = np.max(points[:, 0])
+
+    # compute bounds of interpolation area
+    if x_min_bound is None:
+        x_min_bound = x_min
+    if x_max_bound is None:
+        x_max_bound = x_max
+
+    # explicit formula for quadratic interpolation
+    if no_points == 2 and order == 2 and plot is False:
+        # Solution to quadratic interpolation is given by:
+        # a = -(f1 - f2 - g1(x1 - x2))/(x1 - x2)^2
+        # x_min = x1 - g1/(2a)
+        # if x1 = 0, then is given by:
+        # x_min = - (g1*x2^2)/(2(f2 - f1 - g1*x2))
+
+        if points[0, 0] == 0:
+            x_sol = -points[0, 2] * points[1, 0] ** 2 / (2 * (points[1, 1] - points[0, 1] - points[0, 2] * points[1, 0]))
+        else:
+            a = -(points[0, 1] - points[1, 1] - points[0, 2] * (points[0, 0] - points[1, 0])) / (points[0, 0] - points[1, 0]) ** 2
+            x_sol = points[0, 0] - points[0, 2]/(2*a)
+
+        x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+
+    # explicit formula for cubic interpolation
+    elif no_points == 2 and order == 3 and plot is False:
+        # Solution to cubic interpolation is given by:
+        # d1 = g1 + g2 - 3((f1 - f2)/(x1 - x2))
+        # d2 = sqrt(d1^2 - g1*g2)
+        # x_min = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2))
+        d1 = points[0, 2] + points[1, 2] - 3 * ((points[0, 1] - points[1, 1]) / (points[0, 0] - points[1, 0]))
+        value = d1 ** 2 - points[0, 2] * points[1, 2]
+        if value > 0:
+            d2 = np.sqrt(value)
+            x_sol = points[1, 0] - (points[1, 0] - points[0, 0]) * ((points[1, 2] + d2 - d1) / (points[1, 2] - points[0, 2] + 2 * d2))
+            x_sol = np.minimum(np.maximum(x_min_bound, x_sol), x_max_bound)
+        else:
+            x_sol = (x_max_bound + x_min_bound)/2
+
+    # solve linear system
+    else:
+        # define linear constraints
+        A = np.zeros((0, order + 1))
+        b = np.zeros((0, 1))
+
+        # add linear constraints on function values
+        for i in range(no_points):
+            if not np.isnan(points[i, 1]):
+                constraint = np.zeros((1, order + 1))
+                for j in range(order, -1, -1):
+                    constraint[0, order - j] = points[i, 0] ** j
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 1])
+
+        # add linear constraints on gradient values
+        for i in range(no_points):
+            if not np.isnan(points[i, 2]):
+                constraint = np.zeros((1, order + 1))
+                for j in range(order):
+                    constraint[0, j] = (order - j) * points[i, 0] ** (order - j - 1)
+                A = np.append(A, constraint, 0)
+                b = np.append(b, points[i, 2])
+
+        # check if system is solvable
+        if A.shape[0] != A.shape[1] or np.linalg.matrix_rank(A) != A.shape[0]:
+            x_sol = (x_min_bound + x_max_bound)/2
+            f_min = np.inf
+        else:
+            # solve linear system for interpolating polynomial
+            coeff = np.linalg.solve(A, b)
+
+            # compute critical points
+            dcoeff = np.zeros(order)
+            for i in range(len(coeff) - 1):
+                dcoeff[i] = coeff[i] * (order - i)
+
+            crit_pts = np.array([x_min_bound, x_max_bound])
+            crit_pts = np.append(crit_pts, points[:, 0])
+
+            if not np.isinf(dcoeff).any():
+                roots = np.roots(dcoeff)
+                crit_pts = np.append(crit_pts, roots)
+
+            # test critical points
+            f_min = np.inf
+            x_sol = (x_min_bound + x_max_bound) / 2 # defaults to bisection
+            for crit_pt in crit_pts:
+                if np.isreal(crit_pt):
+                    if not np.isrealobj(crit_pt): crit_pt = crit_pt.real
+                    if crit_pt >= x_min_bound and crit_pt <= x_max_bound:
+                        F_cp = np.polyval(coeff, crit_pt)
+                        if np.isreal(F_cp) and F_cp < f_min:
+                            x_sol = np.real(crit_pt)
+                            f_min = np.real(F_cp)
+
+            if(plot):
+                import matplotlib.pyplot as plt
+                plt.figure()
+                x = np.arange(x_min_bound, x_max_bound, (x_max_bound - x_min_bound)/10000)
+                f = np.polyval(coeff, x)
+                plt.plot(x, f)
+                plt.plot(x_sol, f_min, 'x')
+
+    return x_sol
+
+
+# polynomial interpolation
+# this code is based on https://github.com/hjmshi/PyTorch-LBFGS/blob/master/functions/LBFGS.py
+# PyTorch-LBFGS: A PyTorch Implementation of L-BFGS
+# this one is modified where instead of clipping the solution by bounds, it tries a lower degree polynomial
+# all the way to bisection
+def _within_bounds(x, lb, ub):
+    if lb is not None and x < lb: return False
+    if ub is not None and x > ub: return False
+    return True
+
+def _quad_interp(points):
+    assert points.shape[0] == 2, points.shape
+    if points[0, 0] == 0:
+        denom = 2 * (points[1, 1] - points[0, 1] - points[0, 2] * points[1, 0])
+        if abs(denom) > 1e-32:
+            return -points[0, 2] * points[1, 0] ** 2 / denom
+    else:
+        denom = (points[0, 0] - points[1, 0]) ** 2
+        if denom > 1e-32:
+            a = -(points[0, 1] - points[1, 1] - points[0, 2] * (points[0, 0] - points[1, 0])) / denom
+            if a > 1e-32:
+                return points[0, 0] - points[0, 2]/(2*a)
+    return None
+
+def _cubic_interp(points, lb, ub):
+    assert points.shape[0] == 2, points.shape
+    denom = points[0, 0] - points[1, 0]
+    if abs(denom) > 1e-32:
+        d1 = points[0, 2] + points[1, 2] - 3 * ((points[0, 1] - points[1, 1]) / denom)
+        value = d1 ** 2 - points[0, 2] * points[1, 2]
+        if value > 0:
+            d2 = np.sqrt(value)
+            denom = points[1, 2] - points[0, 2] + 2 * d2
+            if abs(denom) > 1e-32:
+                x_sol = points[1, 0] - (points[1, 0] - points[0, 0]) * ((points[1, 2] + d2 - d1) / denom)
+                if _within_bounds(x_sol, lb, ub): return x_sol
+
+    # try quadratic interpolations
+    x_sol = _quad_interp(points)
+    if x_sol is not None and _within_bounds(x_sol, lb, ub): return x_sol
+
+    return None
+
+def _poly_interp(points, lb, ub):
+    no_points = points.shape[0]
+    assert no_points > 2, points.shape
+    order = np.sum(1 - np.isnan(points[:, 1:3]).astype('int')) - 1
+
+    # define linear constraints
+    A = np.zeros((0, order + 1))
+    b = np.zeros((0, 1))
+
+    # add linear constraints on function values
+    for i in range(no_points):
+        if not np.isnan(points[i, 1]):
+            constraint = np.zeros((1, order + 1))
+            for j in range(order, -1, -1):
+                constraint[0, order - j] = points[i, 0] ** j
+            A = np.append(A, constraint, 0)
+            b = np.append(b, points[i, 1])
+
+    # add linear constraints on gradient values
+    for i in range(no_points):
+        if not np.isnan(points[i, 2]):
+            constraint = np.zeros((1, order + 1))
+            for j in range(order):
+                constraint[0, j] = (order - j) * points[i, 0] ** (order - j - 1)
+            A = np.append(A, constraint, 0)
+            b = np.append(b, points[i, 2])
+
+    # check if system is solvable
+    if A.shape[0] != A.shape[1] or np.linalg.matrix_rank(A) != A.shape[0]:
+        return None
+
+    # solve linear system for interpolating polynomial
+    coeff = np.linalg.solve(A, b)
+
+    # compute critical points
+    dcoeff = np.zeros(order)
+    for i in range(len(coeff) - 1):
+        dcoeff[i] = coeff[i] * (order - i)
+
+    lower = np.min(points[:, 0]) if lb is None else lb
+    upper = np.max(points[:, 0]) if ub is None else ub
+
+    crit_pts = np.array([lower, upper])
+    crit_pts = np.append(crit_pts, points[:, 0])
+
+    if not np.isinf(dcoeff).any():
+        roots = np.roots(dcoeff)
+        crit_pts = np.append(crit_pts, roots)
+
+    # test critical points
+    f_min = np.inf
+    x_sol = None
+    for crit_pt in crit_pts:
+        if np.isreal(crit_pt):
+            if not np.isrealobj(crit_pt): crit_pt = crit_pt.real
+            if _within_bounds(crit_pt, lb, ub):
+                F_cp = np.polyval(coeff, crit_pt)
+                if np.isreal(F_cp) and F_cp < f_min:
+                    x_sol = np.real(crit_pt)
+                    f_min = np.real(F_cp)
+
+    return x_sol
+
+def polyinterp2(points, lb, ub, unbounded: bool = False):
+    no_points = points.shape[0]
+    if no_points <= 1:
+        return (lb + ub)/2
+
+    order = np.sum(1 - np.isnan(points[:, 1:3]).astype('int')) - 1
+
+    x_min = np.min(points[:, 0])
+    x_max = np.max(points[:, 0])
+
+    # compute bounds of interpolation area
+    if not unbounded:
+        if lb is None:
+            lb = x_min
+        if ub is None:
+            ub = x_max
+
+    if no_points == 2 and order == 2:
+        x_sol = _quad_interp(points)
+        if x_sol is not None and _within_bounds(x_sol, lb, ub): return x_sol
+        return (lb + ub)/2
+
+    if no_points == 2 and order == 3:
+        x_sol = _cubic_interp(points, lb, ub) # includes fallback on _quad_interp
+        if x_sol is not None and _within_bounds(x_sol, lb, ub): return x_sol
+        return (lb + ub)/2
+
+    if no_points <= 2: # order < 2
+        return (lb + ub)/2
+
+    if no_points == 3:
+        for p in (points[:2], points[1:], points[::2]):
+            x_sol = _cubic_interp(p, lb, ub)
+            if x_sol is not None and _within_bounds(x_sol, lb, ub): return x_sol
+
+    x_sol = _poly_interp(points, lb, ub)
+    if x_sol is not None and _within_bounds(x_sol, lb, ub): return x_sol
+    return polyinterp2(points[1:], lb, ub)