torchzero 0.1.8__py3-none-any.whl → 0.3.2__py3-none-any.whl

torchzero/modules/gradient_approximation/rfdm.py

@@ -1,125 +0,0 @@
-from typing import Literal, Any, cast
-
-import torch
-
-from ...utils.python_tools import _ScalarLoss
-from ...tensorlist import Distributions, TensorList
-from ...core import _ClosureType, OptimizerModule, OptimizationVars
-from ._fd_formulas import _FD_Formulas
-from .base_approximator import GradientApproximatorBase
-
-def _two_point_rcd_(closure: _ClosureType, params: TensorList, perturbation: TensorList, eps: TensorList, fx0: _ScalarLoss | None, ):
-    """Two point randomized finite difference (same signature for all other finite differences functions).
-
-    Args:
-        closure (Callable): A closure that reevaluates the model and returns the loss.
-        params (TensorList): TensorList with parameters.
-        perturbation (TensorList): TensorList with perturbation ALREADY MULTIPLIED BY EPSILON.
-        eps (TensorList): Finite difference epsilon.
-        fx0 (ScalarType): Loss at fx0, to avoid reevaluating it each time. On some functions can be None when it isn't needed.
-
-    Returns:
-        TensorList with gradient estimation and approximate loss.
-    """
-    # positive loss
-    params += perturbation
-    loss_pos = closure(False)
-
-    # negative loss
-    params.sub_(perturbation, alpha = 2)
-    loss_neg = closure(False)
-
-    # restore params
-    params += perturbation
-
-    # calculate gradient estimation using central finite differences formula
-    # (we square eps in denominator because perturbation is already multiplied by eps)
-    # grad_est = (perturbation * (loss_pos - loss_neg)) / (2 * eps**2)
-    # is equivalent to the following:
-    return perturbation * eps.map(lambda x: (loss_pos - loss_neg) / (2 * x**2)), loss_pos
-    # also we can't reuse the perturbatuion tensor and multiply it in place,
-    # since if randomize_every is more than 1, that would break it.
-
-def _two_point_rfd_(closure: _ClosureType, params: TensorList, perturbation: TensorList, eps: TensorList, fx0: _ScalarLoss | None):
-    if fx0 is None: raise ValueError()
-
-    params += perturbation
-    fx1 = closure(False)
-
-    params -= perturbation
-
-    return perturbation * eps.map(lambda x: (fx1 - fx0) / x**2), fx0
-
-def _two_point_rbd_(closure: _ClosureType, params: TensorList, perturbation: TensorList, eps: TensorList, fx0: _ScalarLoss | None):
-    if fx0 is None: raise ValueError()
-
-    params -= perturbation
-    fx1 = closure(False)
-
-    params += perturbation
-
-    return perturbation * eps.map(lambda x: (fx0 - fx1) / x**2), fx0
-
-
-class RandomizedFDM(GradientApproximatorBase):
-    """Gradient approximation via randomized finite difference.
-
-    Args:
-        eps (float, optional): finite difference epsilon. Defaults to 1e-5.
-        formula (_FD_Formulas, optional): Finite difference formula. Defaults to 'forward'.
-        n_samples (int, optional): number of times gradient is approximated and then averaged. Defaults to 1.
-        distribution (Distributions, optional): distribution for random perturbations. Defaults to "normal".
-        target (str, optional):
-            determines what this module sets.
-
-            "ascent" - it creates a new ascent direction but doesn't treat is as gradient.
-
-            "grad" - it creates the gradient and sets it to `.grad` attributes (default).
-
-            "closure" - it makes a new closure that sets the estimated gradient to the `.grad` attributes.
-    """
-    def __init__(
-        self,
-        eps: float = 1e-5,
-        formula: _FD_Formulas = "forward",
-        n_samples: int = 1,
-        distribution: Distributions = "normal",
-        target: Literal['ascent', 'grad', 'closure'] = 'grad',
-    ):
-        defaults = dict(eps = eps)
-
-        if formula == 'forward':
-            self._finite_difference = _two_point_rfd_
-            requires_fx0 = True
-
-        elif formula == 'backward':
-            self._finite_difference = _two_point_rbd_
-            requires_fx0 = True
-
-        elif formula == 'central':
-            self._finite_difference = _two_point_rcd_
-            requires_fx0 = False
-
-        else: raise ValueError(f"Unknown formula: {formula}")
-
-        self.n_samples = n_samples
-        self.distribution: Distributions = distribution
-
-        super().__init__(defaults, requires_fx0=requires_fx0, target = target)
-
-    @torch.no_grad
-    def _make_ascent(self, closure, params, fx0):
-        eps = self.get_group_key('eps')
-        fx0_approx = None
-
-        if self.n_samples == 1:
-            grads, fx0_approx = self._finite_difference(closure, params, params.sample_like(eps, self.distribution), eps, fx0)
-
-        else:
-            grads = params.zeros_like()
-            for i in range(self.n_samples):
-                g, fx0_approx = self._finite_difference(closure, params, params.sample_like(eps, self.distribution), eps, fx0)
-                grads += g
-            grads /= self.n_samples
-
-        return grads, fx0, fx0_approx

torchzero/modules/line_search/armijo.py

@@ -1,56 +0,0 @@
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizationVars
-from .base_ls import LineSearchBase
-
-
-class ArmijoLS(LineSearchBase):
-    """Armijo backtracking line search
-
-    Args:
-        alpha (float): initial step size.
-        mul (float, optional): lr multiplier on each iteration. Defaults to 0.5.
-        beta (float, optional):
-            armijo condition parameter, fraction of expected linear loss decrease to accept.
-            Larger values mean loss needs to decrease more for a step sizer to be accepted. Defaults to 1e-4.
-        max_iter (int, optional): maximum iterations. Defaults to 10.
-        log_lrs (bool, optional): logs learning rates. Defaults to False.
-    """
-    def __init__(
-        self,
-        alpha: float = 1,
-        mul: float = 0.5,
-        beta: float = 1e-2,
-        max_iter: int = 10,
-        log_lrs = False,
-    ):
-        defaults = dict(alpha=alpha)
-        super().__init__(defaults, maxiter=None, log_lrs=log_lrs)
-        self.mul = mul
-        self.beta = beta
-        self.max_iter = max_iter
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        if vars.closure is None: raise RuntimeError(f"Line searches ({self.__class__.__name__}) require a closure")
-        ascent = vars.maybe_use_grad_(params)
-        grad = vars.maybe_compute_grad_(params)
-        alpha = self.get_first_group_key('alpha')
-        if vars.fx0 is None: vars.fx0 = vars.closure(False)
-
-        # loss decrease per lr=1 if function was linear
-        decrease_per_lr = (grad*ascent).total_sum()
-
-        for _ in range(self.max_iter):
-            loss = self._evaluate_lr_(alpha, vars.closure, ascent, params)
-
-            # expected decrease
-            expected_decrease = decrease_per_lr * alpha
-
-            if (vars.fx0 - loss) / expected_decrease >= self.beta:
-                return alpha
-
-            alpha *= self.mul
-
-        return 0

torchzero/modules/line_search/base_ls.py

@@ -1,139 +0,0 @@
-from typing import Literal
-from abc import ABC, abstractmethod
-
-
-import torch
-
-from ...tensorlist import TensorList
-from ...core import _ClosureType, OptimizationVars, OptimizerModule, _maybe_pass_backward
-from ...utils.python_tools import _ScalarLoss
-
-
-class MaxIterReached(Exception): pass
-
-class LineSearchBase(OptimizerModule, ABC):
-    """Base linesearch class. This is an abstract class, please don't use it as the optimizer.
-
-    When inheriting from this class the easiest way is only override `_find_best_lr`, which should
-    return the final lr to use.
-
-    Args:
-        defaults (dict): dictionary with default parameters for the module.
-        target (str, optional):
-            determines how _update method is used in the default step method.
-
-            "ascent" - it updates the ascent
-
-            "grad" - it updates the gradient (and sets `.grad` attributes to updated gradient).
-
-            "closure" - it makes a new closure that sets the updated ascent to the .`grad` attributes.
-        maxiter (_type_, optional): maximum line search iterations
-            (useful for things like scipy.optimize.minimize_scalar) as it doesn't have
-            an exact iteration limit. Defaults to None.
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-    """
-    def __init__(
-        self,
-        defaults: dict,
-        target: Literal['grad', 'ascent', 'closure'] = 'ascent',
-        maxiter=None,
-        log_lrs=False,
-    ):
-        super().__init__(defaults, target=target)
-        self._reset()
-
-        self.maxiter = maxiter
-        self.log_lrs = log_lrs
-        self._lrs: list[dict[float, _ScalarLoss]] = []
-        """this only gets filled if `log_lrs` is True. On each step, a dictionary is added to this list,
-        with all lrs tested at that step as keys and corresponding losses as values."""
-
-    def _reset(self):
-        """Resets `_last_lr`, `_lowest_loss`, `_best_lr`, `_fx0_approx` and `_current_iter`."""
-        self._last_lr = 0
-        self._lowest_loss = float('inf')
-        self._best_lr = 0
-        self._fx0_approx = None
-        self._current_iter = 0
-
-    def _set_lr_(self, lr: float, ascent_direction: TensorList, params: TensorList, ):
-        alpha = self._last_lr - lr
-        if alpha != 0: params.add_(ascent_direction, alpha = alpha)
-        self._last_lr = lr
-
-    # lr is first here so that we can use a partial
-    def _evaluate_lr_(self, lr: float, closure: _ClosureType, ascent: TensorList, params: TensorList, backward=False):
-        """Evaluate `lr`, if loss is better than current lowest loss,
-        overrides `self._lowest_loss` and `self._best_lr`.
-
-        Args:
-            closure (ClosureType): closure.
-            params (tl.TensorList): params.
-            ascent_direction (tl.TensorList): ascent.
-            lr (float): lr to evaluate.
-
-        Returns:
-            Loss with evaluated lr.
-        """
-        # check max iter
-        if self._current_iter == self.maxiter: raise MaxIterReached
-        self._current_iter += 1
-
-        # set new lr and evaluate loss with it
-        self._set_lr_(lr, ascent, params = params)
-        with torch.enable_grad() if backward else torch.no_grad(): self._fx0_approx = _maybe_pass_backward(closure, backward)
-
-        # if it is the best so far, record it
-        if self._fx0_approx < self._lowest_loss:
-            self._lowest_loss = self._fx0_approx
-            self._best_lr = lr
-
-        # log lr and loss
-        if self.log_lrs:
-            self._lrs[-1][lr] = self._fx0_approx
-
-        return self._fx0_approx
-
-    def _evaluate_lr_ensure_float(
-        self,
-        lr: float,
-        closure: _ClosureType,
-        ascent: TensorList,
-        params: TensorList,
-    ) -> float:
-        """Same as _evaluate_lr_ but ensures that the loss value is float."""
-        v = self._evaluate_lr_(lr, closure, ascent, params)
-        if isinstance(v, torch.Tensor): return v.detach().cpu().item()
-        return float(v)
-
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        """This should return the best lr."""
-        ... # pylint:disable=unnecessary-ellipsis
-
-    @torch.no_grad
-    def step(self, vars: OptimizationVars):
-        self._reset()
-        if self.log_lrs: self._lrs.append({})
-
-        params = self.get_params()
-        ascent_direction = vars.maybe_use_grad_(params)
-
-        try:
-            lr = self._find_best_lr(vars, params) # pylint:disable=assignment-from-no-return
-        except MaxIterReached:
-            lr = self._best_lr
-
-        # if child is None, set best lr which update params and return loss
-        if self.next_module is None:
-            self._set_lr_(lr, ascent_direction, params)
-            return self._lowest_loss
-
-        # otherwise undo the update by setting lr to 0 and instead multiply ascent direction by lr.
-        self._set_lr_(0, ascent_direction, params)
-        ascent_direction.mul_(self._best_lr)
-        vars.ascent = ascent_direction
-        if vars.fx0_approx is None: vars.fx0_approx = self._lowest_loss
-        return self.next_module.step(vars)
-

torchzero/modules/line_search/directional_newton.py

@@ -1,217 +0,0 @@
-import numpy as np
-import torch
-
-from ...tensorlist import TensorList
-from ...core import OptimizationVars
-from .base_ls import LineSearchBase
-
-_FloatOrTensor = float | torch.Tensor
-def _fit_and_minimize_quadratic_2points_grad(x1:_FloatOrTensor,y1:_FloatOrTensor,y1_prime:_FloatOrTensor,x2:_FloatOrTensor,y2:_FloatOrTensor):
-    """Fits a quadratic to value and gradient and x1 and value at x2 and returns minima and a parameter."""
-    a = (y1_prime * x2 - y2 - y1_prime*x1 + y1) / (x1**2 - x2**2 - 2*x1**2 + 2*x1*x2)
-    b = y1_prime - 2*a*x1
-    # c = -(a*x1**2 + b*x1 - y1)
-    return -b / (2 * a), a
-
-def _ensure_float(x):
-    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
-    elif isinstance(x, np.ndarray): return x.item()
-    return float(x)
-
-class DirectionalNewton(LineSearchBase):
-    """Minimizes a parabola in the direction of the update via one additional forward pass,
-    and uses another forward pass to make sure it didn't overstep (optionally).
-    So in total this performs three forward passes and one backward.
-
-    It is recommented to put LR before DirectionalNewton.
-
-    First forward and backward pass is used to calculate the value and gradient at initial parameters.
-    Then a gradient descent step is performed with `lr` learning rate, and loss is recalculated
-    with new parameters. A quadratic is fitted to two points and gradient,
-    if it has positive curvature, this makes a step towards the minimum, and checks if lr decreased
-    with an additional forward pass.
-
-    Args:
-        eps (float, optional):
-            learning rate, also functions as epsilon for directional second derivative estimation. Defaults to 1.
-        max_dist (float | None, optional):
-            maximum distance to step when minimizing quadratic.
-            If minimum is further than this distance, minimization is not performed. Defaults to 1e4.
-        validate_step (bool, optional):
-            uses an additional forward pass to check
-            if step towards the minimum actually decreased the loss. Defaults to True.
-        alpha (float, optional):
-            epsilon for estimating directional second derivative, also works as learning rate
-            for when curvature is negative or loss increases.
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-
-    Note:
-        While lr scheduling is supported, this uses lr of the first parameter for all parameters.
-    """
-    def __init__(self, max_dist: float | None = 1e5, validate_step = True, alpha:float=1, log_lrs = False,):
-        super().__init__({"alpha": alpha}, maxiter=None, log_lrs=log_lrs)
-
-        self.max_dist = max_dist
-        self.validate_step = validate_step
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        if vars.closure is None: raise ValueError('QuardaticLS requires closure')
-        closure = vars.closure
-
-        params = self.get_params()
-        grad = vars.maybe_compute_grad_(params)
-        ascent = vars.maybe_use_grad_(params)
-        if vars.fx0 is None: vars.fx0 = vars.closure(False) # at this stage maybe_compute_grad could've evaluated fx0
-
-        alpha: float = self.get_first_group_key('alpha') # this doesn't support variable lrs but we still want to support schedulers
-
-        # directional f'(x1)
-        y1_prime = (grad * ascent).total_sum()
-
-        # f(x2)
-        y2 = self._evaluate_lr_(alpha, closure, ascent, params)
-
-        # if gradients weren't 0
-        if y1_prime != 0:
-            xmin, a = _fit_and_minimize_quadratic_2points_grad(
-                x1=0,
-                y1=vars.fx0,
-                y1_prime=-y1_prime,
-                x2=alpha,
-                # we stepped in the direction of minus gradient times lr.
-                # which is why y1_prime is negative and we multiply x2 by lr.
-                y2=y2
-            )
-            # so we obtained xmin in lr*grad units. We need in lr units.
-            xmin = _ensure_float(xmin)
-
-            # make sure curvature is positive
-            if a > 0:
-
-                # discard very large steps
-                if self.max_dist is None or xmin <= self.max_dist:
-
-                    # if validate_step is enabled, make sure loss didn't increase
-                    if self.validate_step:
-                        y_val = self._evaluate_lr_(xmin, closure, ascent, params)
-                        # if it increased, move back to y2.
-                        if y_val > y2:
-                            return float(alpha)
-
-                    return float(xmin)
-
-        return float(alpha)
-
-def _fit_and_minimize_quadratic_3points(
-    x1: _FloatOrTensor,
-    y1: _FloatOrTensor,
-    x2: _FloatOrTensor,
-    y2: _FloatOrTensor,
-    x3: _FloatOrTensor,
-    y3: _FloatOrTensor,
-):
-    """Fits a quadratic to three points."""
-    a = (x1*(y3-y2) + x2*(y1-y3) + x3*(y2-y1)) / ((x1-x2) * (x1 - x3) * (x2 - x3))
-    b = (y2-y1) / (x2-x1) - a*(x1+x2)
-    # c = (y1 - a*x1**2 - b*x1)
-    return (-b / (2 * a), a)
-
-
-def _newton_step_3points(
-    xneg: _FloatOrTensor,
-    yneg: _FloatOrTensor,
-    x0: _FloatOrTensor,
-    y0: _FloatOrTensor,
-    xpos: _FloatOrTensor, # since points are evenly spaced, xpos is x0 + eps, its turns out unused
-    ypos: _FloatOrTensor,
-):
-    eps = x0 - xneg
-    dx = (-yneg + ypos) / (2 * eps)
-    ddx = (ypos - 2*y0 + yneg) / (eps**2)
-
-    # xneg is actually x0
-    return xneg - dx / ddx, ddx
-
-class DirectionalNewton3Points(LineSearchBase):
-    """Minimizes a parabola in the direction of the update via two additional forward pass,
-    and uses another forward pass to make sure it didn't overstep (optionally).
-    So in total this performs four forward passes.
-
-    It is recommented to put LR before DirectionalNewton3Points
-
-    Two steps are performed in the direction of the update with `lr` learning rate.
-    A quadratic is fitted to three points, if it has positive curvature,
-    this makes a step towards the minimum, and checks if lr decreased
-    with an additional forward pass.
-
-    Args:
-            for when curvature is negative or loss increases.
-        max_dist (float | None, optional):
-            maximum distance to step when minimizing quadratic.
-            If minimum is further than this distance, minimization is not performed. Defaults to 1e4.
-        validate_step (bool, optional):
-            uses an additional forward pass to check
-            if step towards the minimum actually decreased the loss. Defaults to True.
-        alpha (float, optional):
-            epsilon for estimating directional second derivative, also works as learning rate
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-
-    Note:
-        While lr scheduling is supported, this uses lr of the first parameter for all parameters.
-    """
-    def __init__(self, max_dist: float | None = 1e4, validate_step = True, alpha: float = 1, log_lrs = False,):
-        super().__init__(dict(alpha = alpha), maxiter=None, log_lrs=log_lrs)
-
-        self.alpha = alpha
-        self.max_dist = max_dist
-        self.validate_step = validate_step
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        if vars.closure is None: raise ValueError('QuardaticLS requires closure')
-        closure = vars.closure
-        ascent_direction = vars.ascent
-        if ascent_direction is None: raise ValueError('Ascent direction is None')
-        alpha: float = self.get_first_group_key('alpha')
-
-        if vars.fx0 is None: vars.fx0 = vars.closure(False)
-        params = self.get_params()
-
-        # make a step in the direction and evaluate f(x2)
-        y2 = self._evaluate_lr_(alpha, closure, ascent_direction, params)
-
-        # make a step in the direction and evaluate f(x3)
-        y3 = self._evaluate_lr_(alpha*2, closure, ascent_direction, params)
-
-        # if gradients weren't 0
-        xmin, a = _newton_step_3points(
-            0, vars.fx0,
-            # we stepped in the direction of minus ascent_direction.
-            alpha, y2,
-            alpha * 2, y3
-        )
-        xmin = _ensure_float(xmin)
-
-        # make sure curvature is positive
-        if a > 0:
-
-            # discard very large steps
-            if self.max_dist is None or xmin <= self.max_dist:
-
-                # if validate_step is enabled, make sure loss didn't increase
-                if self.validate_step:
-                    y_val = self._evaluate_lr_(xmin, closure, ascent_direction, params)
-                    # if it increased, move back to y2.
-                    if y_val > y2 or y_val > y3:
-                        if y3 > y2: return alpha
-                        else: return alpha * 2
-
-                return xmin
-
-        if y3 > y2: return alpha
-        else: return alpha * 2

torchzero/modules/line_search/grid_ls.py

@@ -1,158 +0,0 @@
-from typing import Any, Literal
-from collections.abc import Sequence
-
-import numpy as np
-import torch
-
-from ...tensorlist import TensorList
-from ...core import _ClosureType, OptimizationVars
-from .base_ls import LineSearchBase
-
-class GridLS(LineSearchBase):
-    """Test all `lrs` and pick best.
-
-    Args:
-        lrs (Sequence[float] | np.ndarray | torch.Tensor): sequence of lrs to test.
-        stop_on_improvement (bool, optional): stops if loss improves compared to current loss. Defaults to False.
-        stop_on_worsened (bool, optional):
-            stops if next lr loss is worse than previous one.
-            this assumes that lrs are in ascending order. Defaults to False.
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-    """
-    def __init__(
-        self,
-        lrs: Sequence[float] | np.ndarray | torch.Tensor,
-        stop_on_improvement=False,
-        stop_on_worsened=False,
-        log_lrs = False,
-    ):
-        super().__init__({}, maxiter=None, log_lrs=log_lrs)
-        self.lrs = lrs
-        self.stop_on_improvement = stop_on_improvement
-        self.stop_on_worsened = stop_on_worsened
-
-    @torch.no_grad
-    def _find_best_lr(self, vars: OptimizationVars, params: TensorList) -> float:
-        if vars.closure is None: raise ValueError("closure is not set")
-        if vars.ascent is None: raise ValueError("ascent_direction is not set")
-
-        if self.stop_on_improvement:
-            if vars.fx0 is None: vars.fx0 = vars.closure(False)
-            self._lowest_loss = vars.fx0
-
-        for lr in self.lrs:
-            loss = self._evaluate_lr_(float(lr), vars.closure, vars.ascent, params)
-
-            # if worsened
-            if self.stop_on_worsened and loss != self._lowest_loss:
-                break
-
-            # if improved
-            if self.stop_on_improvement and loss == self._lowest_loss:
-                break
-
-        return float(self._best_lr)
-
-
-
-class MultiplicativeLS(GridLS):
-    """Starts with `init` lr, then keeps multiplying it by `mul` until loss stops decreasing.
-
-    Args:
-        init (float, optional): initial lr. Defaults to 0.001.
-        mul (float, optional): lr multiplier. Defaults to 2.
-        num (int, optional): maximum number of multiplication steps. Defaults to 10.
-        stop_on_improvement (bool, optional): stops if loss improves compared to current loss. Defaults to False.
-        stop_on_worsened (bool, optional):
-            stops if next lr loss is worse than previous one.
-            this assumes that lrs are in ascending order. Defaults to False.
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-    """
-    def __init__(
-        self,
-        init: float = 0.001,
-        mul: float = 2,
-        num=10,
-        stop_on_improvement=False,
-        stop_on_worsened=True,
-    ):
-        super().__init__(
-            [init * mul**i for i in range(num)],
-            stop_on_improvement=stop_on_improvement,
-            stop_on_worsened=stop_on_worsened,
-        )
-
-class BacktrackingLS(GridLS):
-    """tests `init` lr, and keeps multiplying it by `mul` until loss becomes better than initial loss.
-
-    note: this doesn't include Armijo–Goldstein condition.
-
-    Args:
-        init (float, optional): initial lr. Defaults to 1.
-        mul (float, optional): lr multiplier. Defaults to 0.5.
-        num (int, optional): maximum number of multiplication steps. Defaults to 10.
-        stop_on_improvement (bool, optional): stops if loss improves compared to current loss. Defaults to False.
-        stop_on_worsened (bool, optional):
-            stops if next lr loss is worse than previous one.
-            this assumes that lrs are in ascending order. Defaults to False.
-        log_lrs (bool, optional):
-            saves lrs and losses with them into optimizer._lrs (for debugging).
-            Defaults to False.
-
-    """
-    def __init__(
-        self,
-        init: float = 1,
-        mul: float = 0.5,
-        num=10,
-        stop_on_improvement=True,
-        stop_on_worsened=False,
-        log_lrs = False,
-    ):
-        super().__init__(
-            [init * mul**i for i in range(num)],
-            stop_on_improvement=stop_on_improvement,
-            stop_on_worsened=stop_on_worsened,
-            log_lrs = log_lrs,
-        )
-
-class LinspaceLS(GridLS):
-    """Test all learning rates from a linspace and pick best."""
-    def __init__(
-        self,
-        start: float = 0.001,
-        end: float = 2,
-        steps=10,
-        stop_on_improvement=False,
-        stop_on_worsened=False,
-        log_lrs = False,
-    ):
-        super().__init__(
-            torch.linspace(start, end, steps),
-            stop_on_improvement=stop_on_improvement,
-            stop_on_worsened=stop_on_worsened,
-            log_lrs = log_lrs,
-        )
-
-class ArangeLS(GridLS):
-    """Test all learning rates from a linspace and pick best."""
-    def __init__(
-        self,
-        start: float = 0.001,
-        end: float = 2,
-        step=0.1,
-        stop_on_improvement=False,
-        stop_on_worsened=False,
-        log_lrs = False,
-
-    ):
-        super().__init__(
-            torch.arange(start, end, step),
-            stop_on_improvement=stop_on_improvement,
-            stop_on_worsened=stop_on_worsened,
-            log_lrs = log_lrs,
-        )