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

torchzero/modules/line_search/polynomial.py

@@ -1,233 +0,0 @@
-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]))
-        d2 = np.sqrt(d1 ** 2 - points[0, 2] * points[1, 2])
-        if np.isreal(d2):
-            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) and 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
-
-
-
-# class PolynomialLineSearch(LineSearch):
-#     """TODO
-
-#     Line search via polynomial interpolation.
-
-#     Args:
-#         init (float, optional): Initial step size. Defaults to 1.0.
-#         c1 (float, optional): Acceptance value for weak wolfe condition. Defaults to 1e-4.
-#         c2 (float, optional): Acceptance value for strong wolfe condition (set to 0.1 for conjugate gradient). Defaults to 0.9.
-#         maxiter (int, optional): Maximum number of line search iterations. Defaults to 25.
-#         maxzoom (int, optional): Maximum number of zoom iterations. Defaults to 10.
-#         expand (float, optional): Expansion factor (multipler to step size when weak condition not satisfied). Defaults to 2.0.
-#         adaptive (bool, optional):
-#             when enabled, if line search failed, initial step size is reduced.
-#             Otherwise it is reset to initial value. Defaults to True.
-#         plus_minus (bool, optional):
-#             If enabled and the direction is not descent direction, performs line search in opposite direction. Defaults to False.
-
-
-#     Examples:
-#         Conjugate gradient method with strong wolfe line search. Nocedal, Wright recommend setting c2 to 0.1 for CG.
-
-#         .. code-block:: python
-
-#             opt = tz.Modular(
-#                 model.parameters(),
-#                 tz.m.PolakRibiere(),
-#                 tz.m.StrongWolfe(c2=0.1)
-#             )
-
-#         LBFGS strong wolfe line search:
-
-#         .. code-block:: python
-
-#             opt = tz.Modular(
-#                 model.parameters(),
-#                 tz.m.LBFGS(),
-#                 tz.m.StrongWolfe()
-#             )
-
-#     """
-#     def __init__(
-#         self,
-#         init: float = 1.0,
-#         c1: float = 1e-4,
-#         c2: float = 0.9,
-#         maxiter: int = 25,
-#         maxzoom: int = 10,
-#         # a_max: float = 1e10,
-#         expand: float = 2.0,
-#         adaptive = True,
-#         plus_minus = False,
-#     ):
-#         defaults=dict(init=init,c1=c1,c2=c2,maxiter=maxiter,maxzoom=maxzoom,
-#                       expand=expand, adaptive=adaptive, plus_minus=plus_minus)
-#         super().__init__(defaults=defaults)
-
-#         self.global_state['initial_scale'] = 1.0
-#         self.global_state['beta_scale'] = 1.0
-
-#     @torch.no_grad
-#     def search(self, update, var):
-#         objective = self.make_objective_with_derivative(var=var)
-
-#         init, c1, c2, maxiter, maxzoom, expand, adaptive, plus_minus = itemgetter(
-#             'init', 'c1', 'c2', 'maxiter', 'maxzoom',
-#             'expand', 'adaptive', 'plus_minus')(self.settings[var.params[0]])
-
-#         f_0, g_0 = objective(0)
-
-#         step_size,f_a = strong_wolfe(
-#             objective,
-#             f_0=f_0, g_0=g_0,
-#             init=init * self.global_state.setdefault("initial_scale", 1),
-#             c1=c1,
-#             c2=c2,
-#             maxiter=maxiter,
-#             maxzoom=maxzoom,
-#             expand=expand,
-#             plus_minus=plus_minus,
-#         )
-
-#         if f_a is not None and (f_a > f_0 or _notfinite(f_a)): step_size = None
-#         if step_size is not None and step_size != 0 and not _notfinite(step_size):
-#             self.global_state['initial_scale'] = min(1.0, self.global_state['initial_scale'] * math.sqrt(2))
-#             return step_size
-
-#         # fallback to backtracking on fail
-#         if adaptive: self.global_state['initial_scale'] *= 0.5
-#         return 0

torchzero/modules/momentum/matrix_momentum.py

@@ -1,193 +0,0 @@
-from typing import Literal
-
-import torch
-
-from ...core import Module, apply_transform, Chainable
-from ...utils import NumberList, TensorList, as_tensorlist
-from ...utils.derivatives import hvp, hvp_fd_central, hvp_fd_forward
-
-class MatrixMomentum(Module):
-    """Second order momentum method.
-
-    Matrix momentum is useful for convex objectives, also for some reason it has very really good generalization on elastic net logistic regression.
-
-    .. note::
-        :code:`mu` needs to be tuned very carefully. It is supposed to be smaller than (1/largest eigenvalue), otherwise this will be very unstable.
-
-    .. note::
-        I have devised an adaptive version of this - :code:`tz.m.AdaptiveMatrixMomentum`, and it works well
-        without having to tune :code:`mu`.
-
-    .. note::
-        In most cases MatrixMomentum should be the first module in the chain because it relies on autograd.
-
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating HVPs.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-    Args:
-        mu (float, optional): this has a similar role to (1 - beta) in normal momentum. Defaults to 0.1.
-        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
-        hvp_method (str, optional):
-            Determines how Hessian-vector products are evaluated.
-
-            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
-              This requires creating a graph for the gradient.
-            - ``"forward"``: Use a forward finite difference formula to
-              approximate the HVP. This requires one extra gradient evaluation.
-            - ``"central"``: Use a central finite difference formula for a
-              more accurate HVP approximation. This requires two extra
-              gradient evaluations.
-            Defaults to "autograd".
-        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
-        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
-
-    Reference:
-        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
-    """
-
-    def __init__(
-        self,
-        mu=0.1,
-        beta: float = 1,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
-        h: float = 1e-3,
-        hvp_tfm: Chainable | None = None,
-    ):
-        defaults = dict(mu=mu, beta=beta, hvp_method=hvp_method, h=h)
-        super().__init__(defaults)
-
-        if hvp_tfm is not None:
-            self.set_child('hvp_tfm', hvp_tfm)
-
-    def reset_for_online(self):
-        super().reset_for_online()
-        self.clear_state_keys('prev_update')
-
-    @torch.no_grad
-    def update(self, var):
-        assert var.closure is not None
-        prev_update = self.get_state(var.params, 'prev_update')
-        hvp_method = self.settings[var.params[0]]['hvp_method']
-        h = self.settings[var.params[0]]['h']
-
-        Hvp, _ = self.Hvp(prev_update, at_x0=True, var=var, rgrad=None, hvp_method=hvp_method, h=h, normalize=True, retain_grad=False)
-        Hvp = [t.detach() for t in Hvp]
-
-        if 'hvp_tfm' in self.children:
-            Hvp = TensorList(apply_transform(self.children['hvp_tfm'], Hvp, params=var.params, grads=var.grad, var=var))
-
-        self.store(var.params, "Hvp", Hvp)
-
-
-    @torch.no_grad
-    def apply(self, var):
-        update = TensorList(var.get_update())
-        Hvp, prev_update = self.get_state(var.params, 'Hvp', 'prev_update', cls=TensorList)
-        mu,beta = self.get_settings(var.params, 'mu','beta', cls=NumberList)
-
-        update.add_(prev_update - Hvp*mu)
-        prev_update.set_(update * beta)
-        var.update = update
-        return var
-
-
-class AdaptiveMatrixMomentum(Module):
-    """Second order momentum method.
-
-    Matrix momentum is useful for convex objectives, also for some reason it has very good generalization on elastic net logistic regression.
-
-    .. note::
-        In most cases MatrixMomentum should be the first module in the chain because it relies on autograd.
-
-    .. note::
-        This module requires the a closure passed to the optimizer step,
-        as it needs to re-evaluate the loss and gradients for calculating HVPs.
-        The closure must accept a ``backward`` argument (refer to documentation).
-
-
-    Args:
-        mu_mul (float, optional): multiplier to the estimated mu. Defaults to 1.
-        beta (float, optional): decay for the buffer, this is not part of the original update rule. Defaults to 1.
-        hvp_method (str, optional):
-            Determines how Hessian-vector products are evaluated.
-
-            - ``"autograd"``: Use PyTorch's autograd to calculate exact HVPs.
-              This requires creating a graph for the gradient.
-            - ``"forward"``: Use a forward finite difference formula to
-              approximate the HVP. This requires one extra gradient evaluation.
-            - ``"central"``: Use a central finite difference formula for a
-              more accurate HVP approximation. This requires two extra
-              gradient evaluations.
-            Defaults to "autograd".
-        h (float, optional): finite difference step size if hvp_method is set to finite difference. Defaults to 1e-3.
-        hvp_tfm (Chainable | None, optional): optional module applied to hessian-vector products. Defaults to None.
-
-    Reference:
-        Orr, Genevieve, and Todd Leen. "Using curvature information for fast stochastic search." Advances in neural information processing systems 9 (1996).
-    """
-
-    def __init__(
-        self,
-        mu_mul: float = 1,
-        beta: float = 1,
-        eps=1e-4,
-        hvp_method: Literal["autograd", "forward", "central"] = "autograd",
-        h: float = 1e-3,
-        hvp_tfm: Chainable | None = None,
-    ):
-        defaults = dict(mu_mul=mu_mul, beta=beta, hvp_method=hvp_method, h=h, eps=eps)
-        super().__init__(defaults)
-
-        if hvp_tfm is not None:
-            self.set_child('hvp_tfm', hvp_tfm)
-
-    def reset_for_online(self):
-        super().reset_for_online()
-        self.clear_state_keys('prev_params', 'prev_grad')
-
-    @torch.no_grad
-    def update(self, var):
-        assert var.closure is not None
-        prev_update, prev_params, prev_grad = self.get_state(var.params, 'prev_update', 'prev_params', 'prev_grad', cls=TensorList)
-
-        settings = self.settings[var.params[0]]
-        hvp_method = settings['hvp_method']
-        h = settings['h']
-        eps = settings['eps']
-
-        mu_mul = NumberList(self.settings[p]['mu_mul'] for p in var.params)
-
-        Hvp, _ = self.Hvp(prev_update, at_x0=True, var=var, rgrad=None, hvp_method=hvp_method, h=h, normalize=True, retain_grad=False)
-        Hvp = [t.detach() for t in Hvp]
-
-        if 'hvp_tfm' in self.children:
-            Hvp = TensorList(apply_transform(self.children['hvp_tfm'], Hvp, params=var.params, grads=var.grad, var=var))
-
-        # adaptive part
-        s_k = var.params - prev_params
-        prev_params.copy_(var.params)
-
-        if hvp_method != 'central': assert var.grad is not None
-        grad = var.get_grad()
-        y_k = grad - prev_grad
-        prev_grad.copy_(grad)
-
-        ada_mu = (s_k.global_vector_norm() / (y_k.global_vector_norm() + eps)) * mu_mul
-
-        self.store(var.params, ['Hvp', 'ada_mu'], [Hvp, ada_mu])
-
-    @torch.no_grad
-    def apply(self, var):
-        Hvp, ada_mu = self.get_state(var.params, 'Hvp', 'ada_mu')
-        Hvp = as_tensorlist(Hvp)
-        beta = NumberList(self.settings[p]['beta'] for p in var.params)
-        update = TensorList(var.get_update())
-        prev_update = TensorList(self.state[p]['prev_update'] for p in var.params)
-
-        update.add_(prev_update - Hvp*ada_mu)
-        prev_update.set_(update * beta)
-        var.update = update
-        return var
-

torchzero/modules/optimizers/adagrad.py

@@ -1,165 +0,0 @@
-from operator import itemgetter
-from typing import Literal
-
-import torch
-from ...core import (
-    Chainable,
-    Module,
-    Target,
-    TensorwiseTransform,
-    Transform,
-    Var,
-    apply_transform,
-)
-from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
-from ...utils.linalg import matrix_power_eigh
-from ..functional import add_power_, lerp_power_, root
-
-
-def adagrad_(
-    tensors_: TensorList,
-    sq_sum_: TensorList,
-    alpha: float | NumberList,
-    lr_decay: float | NumberList,
-    eps: float | NumberList,
-    step: int,
-    pow: float = 2,
-    use_sqrt: bool = True,
-    divide: bool = False,
-
-    # inner args
-    inner: Module | None = None,
-    params: list[torch.Tensor] | None = None,
-    grads: list[torch.Tensor] | None = None,
-):
-    """returns `tensors_`"""
-    clr = alpha / (1 + step * lr_decay)
-
-    sq_sum_ = add_power_(tensors_, sum_=sq_sum_, pow=pow)
-
-    if inner is not None:
-        assert params is not None
-        tensors_ = TensorList(apply_transform(inner, tensors_, params=params, grads=grads))
-
-    if divide: sq_sum_ = sq_sum_ / max(step, 1)
-
-    if use_sqrt: tensors_.div_(root(sq_sum_, p=pow, inplace=False).add_(eps)).mul_(clr)
-    else: tensors_.div_(sq_sum_.add(eps)).mul_(clr)
-
-    return tensors_
-
-
-
-class Adagrad(Transform):
-    """Adagrad, divides by sum of past squares of gradients.
-
-    This implementation is identical to :code:`torch.optim.Adagrad`.
-
-    Args:
-        lr_decay (float, optional): learning rate decay. Defaults to 0.
-        initial_accumulator_value (float, optional): initial value of the sum of squares of gradients. Defaults to 0.
-        eps (float, optional): division epsilon. Defaults to 1e-10.
-        alpha (float, optional): step size. Defaults to 1.
-        pow (float, optional): power for gradients and accumulator root. Defaults to 2.
-        use_sqrt (bool, optional): whether to take the root of the accumulator. Defaults to True.
-        inner (Chainable | None, optional): Inner modules that are applied after updating accumulator and before preconditioning. Defaults to None.
-    """
-    def __init__(
-        self,
-        lr_decay: float = 0,
-        initial_accumulator_value: float = 0,
-        eps: float = 1e-10,
-        alpha: float = 1,
-        pow: float = 2,
-        use_sqrt: bool = True,
-        divide: bool=False,
-        inner: Chainable | None = None,
-    ):
-        defaults = dict(alpha = alpha, lr_decay = lr_decay, initial_accumulator_value=initial_accumulator_value,
-                        eps = eps, pow=pow, use_sqrt = use_sqrt, divide=divide)
-        super().__init__(defaults=defaults, uses_grad=False)
-
-        if inner is not None:
-            self.set_child('inner', inner)
-
-    @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
-        tensors = TensorList(tensors)
-        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-
-        lr_decay,alpha,eps = unpack_dicts(settings, 'lr_decay', 'alpha', 'eps', cls=NumberList)
-
-        pow, use_sqrt, divide = itemgetter('pow', 'use_sqrt', 'divide')(settings[0])
-
-        sq_sum = unpack_states(states, tensors, 'sq_sum', cls=TensorList)
-
-        # initialize accumulator on 1st step
-        if step == 1:
-            sq_sum.set_(tensors.full_like([s['initial_accumulator_value'] for s in settings]))
-
-        return adagrad_(
-            tensors,
-            sq_sum_=sq_sum,
-            alpha=alpha,
-            lr_decay=lr_decay,
-            eps=eps,
-            step=self.global_state["step"],
-            pow=pow,
-            use_sqrt=use_sqrt,
-            divide=divide,
-
-            # inner args
-            inner=self.children.get("inner", None),
-            params=params,
-            grads=grads,
-        )
-
-
-
-class FullMatrixAdagrad(TensorwiseTransform):
-    def __init__(self, beta: float | None = None, decay: float | None = None, sqrt:bool=True, concat_params=True, update_freq=1, init: Literal['identity', 'zeros', 'ones', 'GGT'] = 'identity', divide: bool=False, inner: Chainable | None = None):
-        defaults = dict(beta=beta, decay=decay, sqrt=sqrt, init=init, divide=divide)
-        super().__init__(defaults, uses_grad=False, concat_params=concat_params, update_freq=update_freq, inner=inner,)
-
-    @torch.no_grad
-    def update_tensor(self, tensor, param, grad, loss, state, setting):
-        G = tensor.ravel()
-        GG = torch.outer(G, G)
-        decay = setting['decay']
-        beta = setting['beta']
-        init = setting['init']
-
-        if 'GG' not in state:
-            if init == 'identity': state['GG'] = torch.eye(GG.size(0), device=GG.device, dtype=GG.dtype)
-            elif init == 'zeros': state['GG'] =  torch.zeros_like(GG)
-            elif init == 'ones': state['GG'] = torch.ones_like(GG)
-            elif init == 'GGT': state['GG'] = GG.clone()
-            else: raise ValueError(init)
-        if decay is not None: state['GG'].mul_(decay)
-
-        if beta is not None: state['GG'].lerp_(GG, 1-beta)
-        else: state['GG'].add_(GG)
-        state['i'] = state.get('i', 0) + 1 # number of GGTs in sum
-
-    @torch.no_grad
-    def apply_tensor(self, tensor, param, grad, loss, state, setting):
-        GG = state['GG']
-        sqrt = setting['sqrt']
-        divide = setting['divide']
-        if divide: GG = GG/state.get('i', 1)
-
-        if tensor.numel() == 1:
-            GG = GG.squeeze()
-            if sqrt: return tensor / GG.sqrt()
-            return tensor / GG
-
-        try:
-            if sqrt: B = matrix_power_eigh(GG, -1/2)
-            else: return torch.linalg.solve(GG, tensor.ravel()).view_as(tensor) # pylint:disable = not-callable
-
-        except torch.linalg.LinAlgError:
-            scale = 1 / tensor.abs().max()
-            return tensor.mul_(scale.clip(min=torch.finfo(tensor.dtype).eps, max=1)) # conservative scaling
-
-        return (B @ tensor.ravel()).view_as(tensor)
-