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

torchzero/modules/weight_decay/weight_decay.py

@@ -4,7 +4,7 @@ from typing import Literal
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states
+from ...utils import NumberList, TensorList, as_tensorlist, unpack_dicts, unpack_states, Metrics
 
 
 @torch.no_grad
@@ -14,7 +14,7 @@ def weight_decay_(
     weight_decay: float | NumberList,
     ord: int = 2
 ):
-    """returns `grad_`."""
+    """modifies in-place and returns ``grad_``."""
     if ord == 1: return grad_.add_(params.sign().mul_(weight_decay))
     if ord == 2: return grad_.add_(params.mul(weight_decay))
     if ord - 1 % 2 != 0: return grad_.add_(params.pow(ord-1).mul_(weight_decay))
@@ -29,39 +29,38 @@ class WeightDecay(Transform):
         ord (int, optional): order of the penalty, e.g. 1 for L1 and 2 for L2. Defaults to 2.
         target (Target, optional): what to set on var. Defaults to 'update'.
 
-    Examples:
-        Adam with non-decoupled weight decay
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.WeightDecay(1e-3),
-                tz.m.Adam(),
-                tz.m.LR(1e-3)
-            )
-
-        Adam with decoupled weight decay that still scales with learning rate
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Adam(),
-                tz.m.WeightDecay(1e-3),
-                tz.m.LR(1e-3)
-            )
-
-        Adam with fully decoupled weight decay that doesn't scale with learning rate
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Adam(),
-                tz.m.LR(1e-3),
-                tz.m.WeightDecay(1e-6)
-            )
+    ### Examples:
+
+    Adam with non-decoupled weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.WeightDecay(1e-3),
+        tz.m.Adam(),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with decoupled weight decay that still scales with learning rate
+    ```python
+
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.WeightDecay(1e-3),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with fully decoupled weight decay that doesn't scale with learning rate
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.LR(1e-3),
+        tz.m.WeightDecay(1e-6)
+    )
+    ```
 
     """
     def __init__(self, weight_decay: float, ord: int = 2, target: Target = 'update'):
@@ -77,7 +76,7 @@ class WeightDecay(Transform):
         return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay, ord)
 
 class RelativeWeightDecay(Transform):
-    """Weight decay relative to the mean absolute value of update, gradient or parameters depending on value of :code:`norm_input` argument.
+    """Weight decay relative to the mean absolute value of update, gradient or parameters depending on value of ``norm_input`` argument.
 
     Args:
         weight_decay (float): relative weight decay scale.
@@ -85,40 +84,42 @@ class RelativeWeightDecay(Transform):
         norm_input (str, optional):
             determines what should weight decay be relative to. "update", "grad" or "params".
             Defaults to "update".
+        metric (Ords, optional):
+            metric (norm, etc) that weight decay should be relative to.
+            defaults to 'mad' (mean absolute deviation).
         target (Target, optional): what to set on var. Defaults to 'update'.
 
-    Examples:
-        Adam with non-decoupled relative weight decay
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.RelativeWeightDecay(1e-3),
-                tz.m.Adam(),
-                tz.m.LR(1e-3)
-            )
-
-        Adam with decoupled relative weight decay
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Adam(),
-                tz.m.RelativeWeightDecay(1e-3),
-                tz.m.LR(1e-3)
-            )
-
+    ### Examples:
+
+    Adam with non-decoupled relative weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.RelativeWeightDecay(1e-1),
+        tz.m.Adam(),
+        tz.m.LR(1e-3)
+    )
+    ```
+
+    Adam with decoupled relative weight decay
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.RelativeWeightDecay(1e-1),
+        tz.m.LR(1e-3)
+    )
+    ```
     """
     def __init__(
         self,
         weight_decay: float = 0.1,
         ord: int  = 2,
         norm_input: Literal["update", "grad", "params"] = "update",
+        metric: Metrics = 'mad',
         target: Target = "update",
     ):
-        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input)
+        defaults = dict(weight_decay=weight_decay, ord=ord, norm_input=norm_input, metric=metric)
         super().__init__(defaults, uses_grad=norm_input == 'grad', target=target)
 
     @torch.no_grad
@@ -127,6 +128,7 @@ class RelativeWeightDecay(Transform):
 
         ord = settings[0]['ord']
         norm_input = settings[0]['norm_input']
+        metric = settings[0]['metric']
 
         if norm_input == 'update': src = TensorList(tensors)
         elif norm_input == 'grad':
@@ -137,9 +139,8 @@ class RelativeWeightDecay(Transform):
         else:
             raise ValueError(norm_input)
 
-        mean_abs = src.abs().global_mean()
-
-        return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay * mean_abs, ord)
+        norm = src.global_metric(metric)
+        return weight_decay_(as_tensorlist(tensors), as_tensorlist(params), weight_decay * norm, ord)
 
 
 @torch.no_grad
@@ -162,7 +163,7 @@ class DirectWeightDecay(Module):
     @torch.no_grad
     def step(self, var):
         weight_decay = self.get_settings(var.params, 'weight_decay', cls=NumberList)
-        ord = self.settings[var.params[0]]['ord']
+        ord = self.defaults['ord']
 
         decay_weights_(var.params, weight_decay, ord)
         return var

torchzero/modules/zeroth_order/__init__.py

@@ -0,0 +1 @@
+from .cd import CD

torchzero/modules/zeroth_order/cd.py

@@ -0,0 +1,122 @@
+import math
+import random
+import warnings
+from functools import partial
+from typing import Literal
+
+import numpy as np
+import torch
+
+from ...core import Module
+from ...utils import NumberList, TensorList
+
+class CD(Module):
+    """Coordinate descent. Proposes a descent direction along a single coordinate.
+    A line search such as ``tz.m.ScipyMinimizeScalar(maxiter=8)`` or a fixed step size can be used after this.
+
+    Args:
+        h (float, optional): finite difference step size. Defaults to 1e-3.
+        grad (bool, optional):
+            if True, scales direction by gradient estimate. If False, the scale is fixed to 1. Defaults to True.
+        adaptive (bool, optional):
+            whether to adapt finite difference step size, this requires an additional buffer. Defaults to True.
+        index (str, optional):
+            index selection strategy.
+            - "cyclic" - repeatedly cycles through each coordinate, e.g. ``1,2,3,1,2,3,...``.
+            - "cyclic2" - cycles forward and then backward, e.g ``1,2,3,3,2,1,1,2,3,...`` (default).
+            - "random" - picks coordinate randomly.
+        threepoint (bool, optional):
+            whether to use three points (three function evaluatins) to determine descent direction.
+            if False, uses two points, but then ``adaptive`` can't be used. Defaults to True.
+    """
+    def __init__(self, h:float=1e-3, grad:bool=True, adaptive:bool=True, index:Literal['cyclic', 'cyclic2', 'random']="cyclic2", threepoint:bool=True,):
+        defaults = dict(h=h, grad=grad, adaptive=adaptive, index=index, threepoint=threepoint)
+        super().__init__(defaults)
+
+    @torch.no_grad
+    def step(self, var):
+        closure = var.closure
+        if closure is None:
+            raise RuntimeError("CD requires closure")
+
+        params = TensorList(var.params)
+        ndim = params.global_numel()
+
+        grad_step_size = self.defaults['grad']
+        adaptive = self.defaults['adaptive']
+        index_strategy = self.defaults['index']
+        h = self.defaults['h']
+        threepoint = self.defaults['threepoint']
+
+        # ------------------------------ determine index ----------------------------- #
+        if index_strategy == 'cyclic':
+            idx = self.global_state.get('idx', 0) % ndim
+            self.global_state['idx'] = idx + 1
+
+        elif index_strategy == 'cyclic2':
+            idx = self.global_state.get('idx', 0)
+            self.global_state['idx'] = idx + 1
+            if idx >= ndim * 2:
+                idx = self.global_state['idx'] = 0
+            if idx >= ndim:
+                idx  = (2*ndim - idx) - 1
+
+        elif index_strategy == 'random':
+            if 'generator' not in self.global_state:
+                self.global_state['generator'] = random.Random(0)
+            generator = self.global_state['generator']
+            idx = generator.randrange(0, ndim)
+
+        else:
+            raise ValueError(index_strategy)
+
+        # -------------------------- find descent direction -------------------------- #
+        h_vec = None
+        if adaptive:
+            if threepoint:
+                h_vec = self.get_state(params, 'h_vec', init=lambda x: torch.full_like(x, h), cls=TensorList)
+                h = float(h_vec.flat_get(idx))
+            else:
+                warnings.warn("CD adaptive=True only works with threepoint=True")
+
+        f_0 = var.get_loss(False)
+        params.flat_set_lambda_(idx, lambda x: x + h)
+        f_p = closure(False)
+
+        # -------------------------------- threepoint -------------------------------- #
+        if threepoint:
+            params.flat_set_lambda_(idx, lambda x: x - 2*h)
+            f_n = closure(False)
+            params.flat_set_lambda_(idx, lambda x: x + h)
+
+            if adaptive:
+                assert h_vec is not None
+                if f_0 <= f_p and f_0 <= f_n:
+                    h_vec.flat_set_lambda_(idx, lambda x: max(x/2, 1e-10))
+                else:
+                    if abs(f_0 - f_n) < 1e-12 or abs((f_p - f_0) / (f_0 - f_n) - 1) < 1e-2:
+                        h_vec.flat_set_lambda_(idx, lambda x: min(x*2, 1e10))
+
+            if grad_step_size:
+                alpha = (f_p - f_n) / (2*h)
+
+            else:
+                if f_0 < f_p and f_0 < f_n: alpha = 0
+                elif f_p < f_n: alpha = -1
+                else: alpha = 1
+
+        # --------------------------------- twopoint --------------------------------- #
+        else:
+            params.flat_set_lambda_(idx, lambda x: x - h)
+            if grad_step_size:
+                alpha = (f_p - f_0) / h
+            else:
+                if f_p < f_0: alpha = -1
+                else: alpha = 1
+
+        # ----------------------------- create the update ---------------------------- #
+        update = params.zeros_like()
+        update.flat_set_(idx, alpha)
+        var.update = update
+        return var
+

torchzero/optim/root.py

@@ -0,0 +1,65 @@
+"""WIP, untested"""
+from collections.abc import Callable
+
+from abc import abstractmethod
+import torch
+from ..modules.higher_order.multipoint import sixth_order_im1, sixth_order_p6, _solve
+
+def make_evaluate(f: Callable[[torch.Tensor], torch.Tensor]):
+    def evaluate(x, order) -> tuple[torch.Tensor, ...]:
+        """order=0 - returns (f,), order=1 - returns (f, J), order=2 - returns (f, J, H), etc."""
+        n = x.numel()
+
+        if order == 0:
+            f_x = f(x)
+            return (f_x, )
+
+        x.requires_grad_()
+        with torch.enable_grad():
+            f_x = f(x)
+            I = torch.eye(n, device=x.device, dtype=x.dtype),
+            g_x = torch.autograd.grad(f_x, x, I, create_graph=order!=1, is_grads_batched=True)[0]
+            ret = [f_x, g_x]
+            T = g_x
+
+            # get all derivative up to order
+            for o in range(2, order + 1):
+                is_last = o == order
+                I = torch.eye(T.numel(), device=x.device, dtype=x.dtype),
+                T = torch.autograd.grad(T.ravel(), x, I, create_graph=not is_last, is_grads_batched=True)[0]
+                ret.append(T.view(n, n, *T.shape[1:]))
+
+        return tuple(ret)
+
+    return evaluate
+
+class RootBase:
+    @abstractmethod
+    def one_iteration(
+        self,
+        x: torch.Tensor,
+        evaluate: Callable[[torch.Tensor, int], tuple[torch.Tensor, ...]],
+    ) -> torch.Tensor:
+        """"""
+
+
+# ---------------------------------- methods --------------------------------- #
+def newton(x:torch.Tensor, f_j, lstsq:bool=False):
+    f_x, G_x = f_j(x)
+    return x - _solve(G_x, f_x, lstsq=lstsq)
+
+class Newton(RootBase):
+    def __init__(self, lstsq: bool=False): self.lstsq = lstsq
+    def one_iteration(self, x, evaluate): return newton(x, evaluate, self.lstsq)
+
+
+class SixthOrderP6(RootBase):
+    """sixth-order iterative method
+
+    Abro, Hameer Akhtar, and Muhammad Mujtaba Shaikh. "A new time-efficient and convergent nonlinear solver." Applied Mathematics and Computation 355 (2019): 516-536.
+    """
+    def __init__(self, lstsq: bool=False): self.lstsq = lstsq
+    def one_iteration(self, x, evaluate):
+        def f(x): return evaluate(x, 0)[0]
+        def f_j(x): return evaluate(x, 1)
+        return sixth_order_p6(x, f, f_j, self.lstsq)

torchzero/optim/utility/split.py

@@ -11,12 +11,12 @@ class Split(torch.optim.Optimizer):
 
     Example:
 
-    .. code:: py
-
-        opt = Split(
-            torch.optim.Adam(model.encoder.parameters(), lr=0.001),
-            torch.optim.SGD(model.decoder.parameters(), lr=0.1)
-        )
+    ```python
+    opt = Split(
+        torch.optim.Adam(model.encoder.parameters(), lr=0.001),
+        torch.optim.SGD(model.decoder.parameters(), lr=0.1)
+    )
+    ```
     """
     def __init__(self, *optimizers: torch.optim.Optimizer | Iterable[torch.optim.Optimizer]):
         all_params = []
@@ -25,14 +25,14 @@ class Split(torch.optim.Optimizer):
         # gather all params in case user tries to access them from this object
         for i,opt in enumerate(self.optimizers):
             for p in get_params(opt.param_groups, 'all', list):
-                if p not in all_params: all_params.append(p)
+                if id(p) not in [id(pr) for pr in all_params]: all_params.append(p)
                 else: warnings.warn(
                     f'optimizers[{i}] {opt.__class__.__name__} has some duplicate parameters '
                     'that are also in previous optimizers. They will be updated multiple times.')
 
         super().__init__(all_params, {})
 
-    def step(self, closure: Callable | None = None):
+    def step(self, closure: Callable | None = None): # pyright:ignore[reportIncompatibleMethodOverride]
         loss = None
 
         # if closure provided, populate grad, otherwise each optimizer will call closure separately

torchzero/optim/wrappers/directsearch.py

@@ -7,7 +7,6 @@ import numpy as np
 import torch
 from directsearch.ds import DEFAULT_PARAMS
 
-from ...modules.second_order.newton import tikhonov_
 from ...utils import Optimizer, TensorList
 
 

torchzero/optim/wrappers/fcmaes.py

@@ -2,11 +2,12 @@ from collections.abc import Callable
 from functools import partial
 from typing import Any, Literal
 
+import numpy as np
+import torch
+
 import fcmaes
 import fcmaes.optimizer
 import fcmaes.retry
-import numpy as np
-import torch
 
 from ...utils import Optimizer, TensorList
 

torchzero/optim/wrappers/nlopt.py

@@ -75,8 +75,6 @@ class NLOptWrapper(Optimizer):
     so usually you would want to perform a single step, although performing multiple steps will refine the
     solution.
 
-    Some algorithms are buggy with numpy>=2.
-
     Args:
         params: iterable of parameters to optimize or dicts defining parameter groups.
         algorithm (int | _ALGOS_LITERAL): optimization algorithm from https://nlopt.readthedocs.io/en/latest/NLopt_Algorithms/

torchzero/optim/wrappers/optuna.py

@@ -6,7 +6,7 @@ import torch
 
 import optuna
 
-from ...utils import Optimizer
+from ...utils import Optimizer, totensor, tofloat
 
 def silence_optuna():
     optuna.logging.set_verbosity(optuna.logging.WARNING)
@@ -65,6 +65,6 @@ class OptunaSampler(Optimizer):
         params.from_vec_(vec)
 
         loss = closure()
-        with torch.enable_grad(): self.study.tell(trial, loss)
+        with torch.enable_grad(): self.study.tell(trial, tofloat(torch.nan_to_num(totensor(loss), 1e32)))
 
         return loss

torchzero/optim/wrappers/scipy.py

@@ -4,12 +4,17 @@ from functools import partial
 from typing import Any, Literal
 
 import numpy as np
-import scipy.optimize
 import torch
 
+import scipy.optimize
+
 from ...utils import Optimizer, TensorList
-from ...utils.derivatives import jacobian_and_hessian_mat_wrt, jacobian_wrt
-from ...modules.second_order.newton import tikhonov_
+from ...utils.derivatives import (
+    flatten_jacobian,
+    jacobian_and_hessian_mat_wrt,
+    jacobian_wrt,
+)
+
 
 def _ensure_float(x) -> float:
     if isinstance(x, torch.Tensor): return x.detach().cpu().item()
@@ -21,14 +26,6 @@ def _ensure_numpy(x):
     if isinstance(x, np.ndarray): return x
     return np.array(x)
 
-def matrix_clamp(H: torch.Tensor, reg: float):
-    try:
-        eigvals, eigvecs = torch.linalg.eigh(H) # pylint:disable=not-callable
-        eigvals.clamp_(min=reg)
-        return eigvecs @ torch.diag(eigvals) @ eigvecs.mH
-    except Exception:
-        return H
-
 Closure = Callable[[bool], Any]
 
 class ScipyMinimize(Optimizer):
@@ -76,8 +73,6 @@ class ScipyMinimize(Optimizer):
         options = None,
         jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
         hess: Literal['2-point', '3-point', 'cs', 'autograd'] | scipy.optimize.HessianUpdateStrategy = 'autograd',
-        tikhonov: float | None = 0,
-        min_eigval: float | None = None,
     ):
         defaults = dict(lb=lb, ub=ub)
         super().__init__(params, defaults)
@@ -85,12 +80,10 @@ class ScipyMinimize(Optimizer):
         self.constraints = constraints
         self.tol = tol
         self.callback = callback
-        self.min_eigval = min_eigval
         self.options = options
 
         self.jac = jac
         self.hess = hess
-        self.tikhonov: float | None = tikhonov
 
         self.use_jac_autograd = jac.lower() == 'autograd' and (method is None or method.lower() in [
             'cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'slsqp', 'dogleg',
@@ -111,9 +104,7 @@ class ScipyMinimize(Optimizer):
         with torch.enable_grad():
             value = closure(False)
             _, H = jacobian_and_hessian_mat_wrt([value], wrt = params)
-        if self.tikhonov is not None: H = tikhonov_(H, self.tikhonov)
-        if self.min_eigval is not None: H = matrix_clamp(H, self.min_eigval)
-        return H.detach().cpu().numpy()
+        return H.numpy(force=True)
 
     def _objective(self, x: np.ndarray, params: TensorList, closure):
         # set params to x
@@ -122,7 +113,10 @@ class ScipyMinimize(Optimizer):
         # return value and maybe gradients
         if self.use_jac_autograd:
             with torch.enable_grad(): value = _ensure_float(closure())
-            return value, params.ensure_grad_().grad.to_vec().detach().cpu().numpy()
+            grad = params.ensure_grad_().grad.to_vec().numpy(force=True)
+            # slsqp requires float64
+            if self.method.lower() == 'slsqp': grad = grad.astype(np.float64)
+            return value, grad
         return _ensure_float(closure(False))
 
     @torch.no_grad
@@ -135,7 +129,7 @@ class ScipyMinimize(Optimizer):
             else: hess = None
         else: hess = self.hess
 
-        x0 = params.to_vec().detach().cpu().numpy()
+        x0 = params.to_vec().numpy(force=True)
 
         # make bounds
         lb, ub = self.group_vals('lb', 'ub', cls=list)
@@ -167,7 +161,7 @@ class ScipyMinimize(Optimizer):
 
 
 class ScipyRootOptimization(Optimizer):
-    """Optimization via using scipy.root on gradients, mainly for experimenting!
+    """Optimization via using scipy.optimize.root on gradients, mainly for experimenting!
 
     Args:
         params: iterable of parameters to optimize or dicts defining parameter groups.
@@ -248,6 +242,72 @@ class ScipyRootOptimization(Optimizer):
         return res.fun
 
 
+class ScipyLeastSquaresOptimization(Optimizer):
+    """Optimization via using scipy.optimize.least_squares on gradients, mainly for experimenting!
+
+    Args:
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        method (str | None, optional): _description_. Defaults to None.
+        tol (float | None, optional): _description_. Defaults to None.
+        callback (_type_, optional): _description_. Defaults to None.
+        options (_type_, optional): _description_. Defaults to None.
+        jac (T.Literal['2, optional): _description_. Defaults to 'autograd'.
+    """
+    def __init__(
+        self,
+        params,
+        method='trf',
+        jac='autograd',
+        bounds=(-np.inf, np.inf),
+        ftol=1e-8, xtol=1e-8, gtol=1e-8, x_scale=1.0, loss='linear',
+        f_scale=1.0, diff_step=None, tr_solver=None, tr_options=None,
+        jac_sparsity=None, max_nfev=None, verbose=0
+    ):
+        super().__init__(params, {})
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['__class__'], kwargs['jac']
+        self._kwargs = kwargs
+
+        self.jac = jac
+
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
+        # set params to x
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        # return the gradients
+        with torch.enable_grad(): self.value = closure()
+        jac = params.ensure_grad_().grad.to_vec()
+        return jac.numpy(force=True)
+
+    def _hess(self, x: np.ndarray, params: TensorList, closure):
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        with torch.enable_grad():
+            value = closure(False)
+            _, H = jacobian_and_hessian_mat_wrt([value], wrt = params)
+        return H.numpy(force=True)
+
+    @torch.no_grad
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        params = self.get_params()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        if self.jac == 'autograd': jac = partial(self._hess, params = params, closure = closure)
+        else: jac = self.jac
+
+        res = scipy.optimize.least_squares(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            jac=jac, # type:ignore
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.fun
+
+
+
 
 class ScipyDE(Optimizer):
     """Use scipy.minimize.differential_evolution as pytorch optimizer. Note that this performs full minimization on each step,
@@ -510,4 +570,3 @@ class ScipyBrute(Optimizer):
             **self._kwargs
         )
         params.from_vec_(torch.from_numpy(x0).to(device = params[0].device, dtype=params[0].dtype, copy=False))
-        return None