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

torchzero/optim/wrappers/scipy/experimental.py

@@ -0,0 +1,141 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+class ScipyRootOptimization(WrapperBase):
+
+    """Optimization via using scipy.optimize.root 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: Literal[
+            "hybr",
+            "lm",
+            "broyden1",
+            "broyden2",
+            "anderson",
+            "linearmixing",
+            "diagbroyden",
+            "excitingmixing",
+            "krylov",
+            "df-sane",
+        ] = 'hybr',
+        tol: float | None = None,
+        callback = None,
+        options = None,
+        jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
+    ):
+        super().__init__(params, {})
+        self.method = method
+        self.tol = tol
+        self.callback = callback
+        self.options = options
+
+        self.jac = jac
+        if self.jac == 'autograd': self.jac = True
+
+        # those don't require jacobian
+        if self.method.lower() in ('broyden1', 'broyden2', 'anderson', 'linearmixing', 'diagbroyden', 'excitingmixing', 'krylov', 'df-sane'):
+            self.jac = None
+
+    def _objective(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        if self.jac:
+            f, g, H = self._f_g_H(x, params, closure)
+            return g, H
+
+        f, g = self._f_g(x, params, closure)
+        return g
+
+    @torch.no_grad
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+
+        res = scipy.optimize.root(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            method=self.method,
+            tol=self.tol,
+            callback=self.callback,
+            options=self.options,
+            jac = self.jac,
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun
+
+
+class ScipyLeastSquaresOptimization(WrapperBase):
+    """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: list[torch.Tensor], closure):
+        f, g = self._f_g(x, params, closure)
+        return g
+
+    def _hess(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g,H = self._f_g_H(x, params, closure)
+        return H
+
+    @torch.no_grad
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+
+        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.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun
+

torchzero/optim/wrappers/scipy/minimize.py

@@ -0,0 +1,151 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+
+def _use_jac_hess_hessp(method, jac, hess, use_hessp):
+    # those methods can't use hessp
+    if (method is None) or (method.lower() not in ("newton-cg", "trust-ncg", "trust-krylov", "trust-constr")):
+        use_hessp = False
+
+    # those use gradients
+    use_jac_autograd = (jac.lower() == 'autograd') and ((method is None) or (method.lower() in [
+        'cg', 'bfgs', 'newton-cg', 'l-bfgs-b', 'tnc', 'slsqp', 'dogleg',
+        'trust-ncg', 'trust-krylov', 'trust-exact', 'trust-constr',
+    ]))
+
+    # those use hessian/ some of them can use hessp instead
+    use_hess_autograd = (isinstance(hess, str)) and (hess.lower() == 'autograd') and (method is not None) and (method.lower() in [
+        'newton-cg', 'dogleg', 'trust-ncg', 'trust-krylov', 'trust-exact'
+    ])
+
+    # jac in scipy is '2-point', '3-point', 'cs', True or None.
+    if jac == 'autograd':
+        if use_jac_autograd: jac = True
+        else: jac = None
+
+    return jac, use_jac_autograd, use_hess_autograd, use_hessp
+
+class ScipyMinimize(WrapperBase):
+    """Use scipy.minimize.optimize as pytorch optimizer. Note that this performs full minimization on each step,
+    so usually you would want to perform a single step, although performing multiple steps will refine the
+    solution.
+
+    Please refer to https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
+    for a detailed description of args.
+
+    Args:
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        method (str | None, optional): type of solver.
+            If None, scipy will select one of BFGS, L-BFGS-B, SLSQP,
+            depending on whether or not the problem has constraints or bounds.
+            Defaults to None.
+        bounds (optional): bounds on variables. Defaults to None.
+        constraints (tuple, optional): constraints definition. Defaults to ().
+        tol (float | None, optional): Tolerance for termination. Defaults to None.
+        callback (Callable | None, optional): A callable called after each iteration. Defaults to None.
+        options (dict | None, optional): A dictionary of solver options. Defaults to None.
+        jac (str, optional): Method for computing the gradient vector.
+            Only for CG, BFGS, Newton-CG, L-BFGS-B, TNC, SLSQP, dogleg, trust-ncg, trust-krylov, trust-exact and trust-constr.
+            In addition to scipy options, this supports 'autograd', which uses pytorch autograd.
+            This setting is ignored for methods that don't require gradient. Defaults to 'autograd'.
+        hess (str, optional):
+            Method for computing the Hessian matrix.
+            Only for Newton-CG, dogleg, trust-ncg, trust-krylov, trust-exact and trust-constr.
+            This setting is ignored for methods that don't require hessian. Defaults to 'autograd'.
+        tikhonov (float, optional):
+            optional hessian regularizer value. Only has effect for methods that require hessian.
+    """
+    def __init__(
+        self,
+        params,
+        method: Literal['nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg',
+                    'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp',
+                    'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact',
+                    'trust-krylov'] | str | None = None,
+        lb = None,
+        ub = None,
+        constraints = (),
+        tol: float | None = None,
+        callback = None,
+        options = None,
+        jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
+        hess: Literal['2-point', '3-point', 'cs', 'autograd'] | scipy.optimize.HessianUpdateStrategy = 'autograd',
+        use_hessp: bool = True,
+    ):
+        defaults = dict(lb=lb, ub=ub)
+        super().__init__(params, defaults)
+        self.method = method
+        self.constraints = constraints
+        self.tol = tol
+        self.callback = callback
+        self.options = options
+        self.hess = hess
+
+        self.jac, self.use_jac_autograd, self.use_hess_autograd, self.use_hessp = _use_jac_hess_hessp(method, jac, hess, use_hessp)
+
+    def _objective(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        if self.use_jac_autograd:
+            f, g = self._f_g(x, params, closure)
+            if self.method is not None and self.method.lower() == 'slsqp': g = g.astype(np.float64) #  slsqp requires float64
+            return f, g
+
+        return self._f(x, params, closure)
+
+    def _hess(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g,H = self._f_g_H(x, params, closure)
+        return H
+
+    def _hessp(self, x: np.ndarray, p:np.ndarray, params: list[torch.Tensor], closure):
+        f,g,Hvp = self._f_g_Hvp(x, p, params, closure)
+        return Hvp
+
+    @torch.no_grad
+    def step(self, closure: Closure):# pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+
+        # determine hess argument
+        hess = self.hess
+        hessp = None
+        if hess == 'autograd':
+            if self.use_hess_autograd:
+                if self.use_hessp:
+                    hessp = partial(self._hessp, params=params, closure=closure)
+                    hess = None
+                else:
+                    hess = partial(self._hess, params=params, closure=closure)
+            # hess = 'autograd' but method doesn't use hess
+            else:
+                hess = None
+
+
+        if self.method is not None and (self.method.lower() == 'tnc' or self.method.lower() == 'slsqp'):
+            x0 = x0.astype(np.float64) # those methods error without this
+
+        res = scipy.optimize.minimize(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            method=self.method,
+            bounds=bounds,
+            constraints=self.constraints,
+            tol=self.tol,
+            callback=self.callback,
+            options=self.options,
+            jac = self.jac,
+            hess = hess,
+            hessp = hessp
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun

torchzero/optim/wrappers/scipy/sgho.py

@@ -0,0 +1,111 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import scipy.optimize
+import torch
+
+from ....utils import TensorList
+from ..wrapper import WrapperBase
+from .minimize import _use_jac_hess_hessp
+
+Closure = Callable[[bool], Any]
+
+
+class ScipySHGO(WrapperBase):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        constraints = None,
+        n: int = 100,
+        iters: int = 1,
+        callback = None,
+        options: dict | None = None,
+        sampling_method: str = 'simplicial',
+        minimizer_kwargs: dict | None = None,
+        method: Literal['nelder-mead', 'powell', 'cg', 'bfgs', 'newton-cg',
+                    'l-bfgs-b', 'tnc', 'cobyla', 'cobyqa', 'slsqp',
+                    'trust-constr', 'dogleg', 'trust-ncg', 'trust-exact',
+                    'trust-krylov'] | str = 'l-bfgs-b',
+        jac: Literal['2-point', '3-point', 'cs', 'autograd'] = 'autograd',
+        hess: Literal['2-point', '3-point', 'cs', 'autograd'] | scipy.optimize.HessianUpdateStrategy = 'autograd',
+        use_hessp: bool = True,
+    ):
+        super().__init__(params, dict(lb=lb, ub=ub))
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__'], kwargs["options"]
+        del kwargs["method"], kwargs["jac"], kwargs["hess"], kwargs["use_hessp"], kwargs["minimizer_kwargs"]
+        self._kwargs = kwargs
+        self.minimizer_kwargs = minimizer_kwargs
+        self.options = options
+        self.method = method
+        self.hess = hess
+
+        self.jac, self.use_jac_autograd, self.use_hess_autograd, self.use_hessp = _use_jac_hess_hessp(method, jac, hess, use_hessp)
+
+
+    def _objective(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        if self.use_jac_autograd:
+            f, g = self._f_g(x, params, closure)
+            if self.method.lower() == 'slsqp': g = g.astype(np.float64) #  slsqp requires float64
+            return f, g
+
+        return self._f(x, params, closure)
+
+    def _hess(self, x: np.ndarray, params: list[torch.Tensor], closure):
+        f,g,H = self._f_g_H(x, params, closure)
+        return H
+
+    def _hessp(self, x: np.ndarray, p:np.ndarray, params: list[torch.Tensor], closure):
+        f,g,Hvp = self._f_g_Hvp(x, p, params, closure)
+        return Hvp
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        bounds = self._get_bounds()
+        assert bounds is not None
+
+        # determine hess argument
+        hess = self.hess
+        hessp = None
+        if hess == 'autograd':
+            if self.use_hess_autograd:
+                if self.use_hessp:
+                    hessp = partial(self._hessp, params=params, closure=closure)
+                    hess = None
+                else:
+                    hess = partial(self._hess, params=params, closure=closure)
+            # hess = 'autograd' but method doesn't use hess
+            else:
+                hess = None
+
+
+        if self.method.lower() in ('tnc', 'slsqp'):
+            x0 = x0.astype(np.float64) # those methods error without this
+
+        minimizer_kwargs = self.minimizer_kwargs.copy() if self.minimizer_kwargs is not None else {}
+        minimizer_kwargs.setdefault("method", self.method)
+
+        options = self.options.copy() if self.options is not None else {}
+        minimizer_kwargs.setdefault("jac", self.jac)
+        minimizer_kwargs.setdefault("hess", hess)
+        minimizer_kwargs.setdefault("hessp", hessp)
+        minimizer_kwargs.setdefault("bounds", bounds)
+
+        res = scipy.optimize.shgo(
+            partial(self._objective, params=params, closure=closure),
+            bounds=bounds,
+            minimizer_kwargs=minimizer_kwargs,
+            options=options,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
+        return res.fun
+

torchzero/optim/wrappers/wrapper.py

@@ -0,0 +1,121 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+import numpy as np
+import torch
+
+from ...utils import TensorList, tonumpy
+from ...utils.derivatives import (
+    flatten_jacobian,
+    jacobian_and_hessian_mat_wrt,
+    jacobian_wrt,
+)
+
+
+class WrapperBase(torch.optim.Optimizer):
+    def __init__(self, params, defaults):
+        super().__init__(params, defaults)
+
+    @torch.no_grad
+    def _f(self, x: np.ndarray, params: list[torch.Tensor], closure) -> float:
+        # set params to x
+        params = TensorList(params)
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        return float(closure(False))
+
+    @torch.no_grad
+    def _fs(self, x: np.ndarray, params: list[torch.Tensor], closure) -> np.ndarray:
+        # set params to x
+        params = TensorList(params)
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        return tonumpy(closure(False)).reshape(-1)
+
+
+    @torch.no_grad
+    def _f_g(self, x: np.ndarray, params: list[torch.Tensor], closure) -> tuple[float, np.ndarray]:
+        # set params to x
+        params = TensorList(params)
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        # compute value and derivatives
+        with torch.enable_grad():
+            value = closure()
+            g = params.grad.fill_none(reference=params).to_vec()
+        return float(value), g.numpy(force=True)
+
+    @torch.no_grad
+    def _f_g_H(self, x: np.ndarray, params: list[torch.Tensor], closure) -> tuple[float, np.ndarray, np.ndarray]:
+        # set params to x
+        params = TensorList(params)
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        # compute value and derivatives
+        with torch.enable_grad():
+            value = closure(False)
+            g, H = jacobian_and_hessian_mat_wrt([value], wrt = params)
+        return float(value), g.numpy(force=True), H.numpy(force=True)
+
+    @torch.no_grad
+    def _f_g_Hvp(self, x: np.ndarray, v: np.ndarray, params: list[torch.Tensor], closure) -> tuple[float, np.ndarray, np.ndarray]:
+        # set params to x
+        params = TensorList(params)
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        # compute value and derivatives
+        with torch.enable_grad():
+            value = closure(False)
+            grad = torch.autograd.grad(value, params, create_graph=True, allow_unused=True, materialize_grads=True)
+            flat_grad = torch.cat([i.reshape(-1) for i in grad])
+            Hvp = torch.autograd.grad(flat_grad, params, torch.as_tensor(v, device=flat_grad.device, dtype=flat_grad.dtype))[0]
+
+        return float(value), flat_grad.numpy(force=True), Hvp.numpy(force=True)
+
+    def _get_params(self) -> list[torch.Tensor]:
+        return [p for g in self.param_groups for p in g["params"]]
+
+    def _get_per_parameter_lb_ub(self):
+        # get per-parameter lb and ub
+        lb = []
+        ub = []
+        for group in self.param_groups:
+            lb.extend([group["lb"]] * len(group["params"]))
+            ub.extend([group["ub"]] * len(group["params"]))
+
+        return lb, ub
+
+    def _get_bounds(self):
+
+        # get per-parameter lb and ub
+        lb, ub = self._get_per_parameter_lb_ub()
+        if all(i is None for i in lb) and all(i is None for i in ub): return None
+
+        params = self._get_params()
+        bounds = []
+        for p, l, u in zip(params, lb, ub):
+            bounds.extend([(l, u)] * p.numel())
+
+        return bounds
+
+    def _get_lb_ub(self, ld:dict | None = None, ud: dict | None = None):
+        if ld is None: ld = {}
+        if ud is None: ud = {}
+
+        # get per-parameter lb and ub
+        lb, ub = self._get_per_parameter_lb_ub()
+
+        params = self._get_params()
+        lb_list = []
+        ub_list = []
+        for p, l, u in zip(params, lb, ub):
+            if l in ld: l = ld[l]
+            if u in ud: l = ud[u]
+            lb_list.extend([l] * p.numel())
+            ub_list.extend([u] * p.numel())
+
+        return lb_list, ub_list
+
+    @abstractmethod
+    def step(self, closure) -> Any: # pyright:ignore[reportIncompatibleMethodOverride] # pylint:disable=signature-differs
+        ...

torchzero/utils/__init__.py

@@ -1,33 +1,15 @@
 from . import tensorlist as tl
-from .compile import (
-    _optional_compiler,
-    benchmark_compile_cpu,
-    benchmark_compile_cuda,
-    enable_compilation,
-    set_compilation,
-)
-from .numberlist import NumberList
-from .optimizer import (
-    Init,
-    ListLike,
-    Optimizer,
-    ParamFilter,
-    get_group_vals,
-    get_params,
-    get_state_vals,
-    unpack_states,
-)
-from .params import (
-    Params,
-    _add_defaults_to_param_groups_,
-    _add_updates_grads_to_param_groups_,
-    _copy_param_groups,
-    _make_param_groups,
-)
+
+from .metrics import evaluate_metric
+from .numberlist import NumberList , maybe_numberlist
+from .optimizer import unpack_states
+
+
 from .python_tools import (
     flatten,
     generic_eq,
     generic_ne,
+    generic_is_none,
     reduce_dim,
     safe_dict_update_,
     unpack_dicts,

torchzero/utils/compile.py

@@ -38,11 +38,11 @@ class _MaybeCompiledFunc:
 _optional_compiler = _OptionalCompiler()
 """this holds .enable attribute, set to True to enable compiling for a few functions that benefit from it."""
 
-def set_compilation(enable: bool=True):
+def enable_compilation(enable: bool=True):
     """`enable` is False by default. When True, certain functions will be compiled, which may not work on some systems like Windows, but it usually improves performance."""
     _optional_compiler.enable = enable
 
-def enable_compilation(fn): return _optional_compiler.enable_compilation(fn)
+def allow_compile(fn): return _optional_compiler.enable_compilation(fn)
 
 def benchmark_compile_cuda(fn, n: int, **kwargs):
     # warmup