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

torchzero/optim/wrappers/nlopt.py

@@ -1,12 +1,11 @@
-from typing import Literal
+from typing import Literal, Any
 from collections.abc import Mapping, Callable
 from functools import partial
 import numpy as np
 import torch
 
 import nlopt
-from ...core import TensorListOptimizer, _ClosureType
-from ...tensorlist import TensorList
+from ...utils import Optimizer, TensorList
 
 _ALGOS_LITERAL = Literal[
     "GN_DIRECT",  # = _nlopt.GN_DIRECT
@@ -56,18 +55,21 @@ _ALGOS_LITERAL = Literal[
 ]
 
 def _ensure_float(x):
-    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
-    if isinstance(x, np.ndarray): return x.item()
+    if isinstance(x, torch.Tensor): return float(x.detach().cpu().item())
+    if isinstance(x, np.ndarray): return float(x.item())
     return float(x)
 
 def _ensure_tensor(x):
-    if isinstance(x, np.ndarray):
-        x.setflags(write=True)
-        return torch.from_numpy(x)
+    try:
+        if isinstance(x, np.ndarray): return torch.as_tensor(x.copy())
+    except SystemError:
+        return None
     return torch.tensor(x, dtype=torch.float32)
 
 inf = float('inf')
-class NLOptOptimizer(TensorListOptimizer):
+Closure = Callable[[bool], Any]
+
+class NLOptOptimizer(Optimizer):
     """Use nlopt as pytorch optimizer, with gradient supplied by pytorch autograd.
     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
@@ -119,8 +121,12 @@ class NLOptOptimizer(TensorListOptimizer):
 
         self._last_loss = None
 
-    def _f(self, x: np.ndarray, grad: np.ndarray, closure: _ClosureType, params: TensorList):
-        params.from_vec_(_ensure_tensor(x).to(params[0], copy=False))
+    def _f(self, x: np.ndarray, grad: np.ndarray, closure, params: TensorList):
+        t = _ensure_tensor(x)
+        if t is None:
+            if self.opt is not None: self.opt.force_stop()
+            return None
+        params.from_vec_(t.to(params[0], copy=False))
         if grad.size > 0:
             with torch.enable_grad(): loss = closure()
             self._last_loss = _ensure_float(loss)
@@ -131,12 +137,11 @@ class NLOptOptimizer(TensorListOptimizer):
         return self._last_loss
 
     @torch.no_grad
-    def step(self, closure: _ClosureType): # pylint: disable = signature-differs
-
+    def step(self, closure: Closure): # pylint: disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
         params = self.get_params()
 
         # make bounds
-        lb, ub = self.get_group_keys('lb', 'ub', cls=list)
+        lb, ub = self.group_vals('lb', 'ub', cls=list)
         lower = []
         upper = []
         for p, l, u in zip(params, lb, ub):
@@ -145,9 +150,10 @@ class NLOptOptimizer(TensorListOptimizer):
             lower.extend([l] * p.numel())
             upper.extend([u] * p.numel())
 
-        x0 = params.to_vec().detach().cpu().numpy()
+        x0 = params.to_vec().detach().cpu().numpy().astype(np.float64)
 
         self.opt = nlopt.opt(self.algorithm, x0.size)
+        self.opt.set_exceptions_enabled(False) # required
         self.opt.set_min_objective(partial(self._f, closure = closure, params = params))
         self.opt.set_lower_bounds(lower)
         self.opt.set_upper_bounds(upper)
@@ -160,6 +166,15 @@ class NLOptOptimizer(TensorListOptimizer):
         if self.xtol_abs is not None: self.opt.set_xtol_abs(self.xtol_abs)
         if self.maxtime is not None: self.opt.set_maxtime(self.maxtime)
 
-        x = self.opt.optimize(x0)
+        self._last_loss = None
+        x = None
+        try:
+            x = self.opt.optimize(x0)
+        except SystemError:
+            pass
+        except Exception as e:
+            raise e from None
+
+        if self._last_loss is None or x is None: return closure(False)
         params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
         return self._last_loss

torchzero/optim/wrappers/scipy.py

@@ -1,20 +1,15 @@
-from typing import Literal, Any
 from collections import abc
+from collections.abc import Callable
 from functools import partial
+from typing import Any, Literal
 
 import numpy as np
-import torch
-
 import scipy.optimize
+import torch
 
-from ...core import _ClosureType, TensorListOptimizer
-from ...utils.derivatives import jacobian, jacobian_list_to_vec, hessian, hessian_list_to_mat, jacobian_and_hessian
-from ...modules import WrapClosure
-from ...modules.experimental.subspace import Projection, Proj2Masks, ProjGrad, ProjNormalize, Subspace
-from ...modules.second_order.newton import regularize_hessian_
-from ...tensorlist import TensorList
-from ..modular import Modular
-
+from ...utils import Optimizer, TensorList
+from ...utils.derivatives import jacobian_and_hessian_mat_wrt, jacobian_wrt
+from ...modules.second_order.newton import tikhonov_
 
 def _ensure_float(x):
     if isinstance(x, torch.Tensor): return x.detach().cpu().item()
@@ -26,7 +21,17 @@ def _ensure_numpy(x):
     if isinstance(x, np.ndarray): return x
     return np.array(x)
 
-class ScipyMinimize(TensorListOptimizer):
+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):
     """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.
@@ -71,7 +76,8 @@ class ScipyMinimize(TensorListOptimizer):
         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 | Literal['eig'] = 0,
+        tikhonov: float | None = 0,
+        min_eigval: float | None = None,
     ):
         defaults = dict(lb=lb, ub=ub)
         super().__init__(params, defaults)
@@ -79,11 +85,12 @@ class ScipyMinimize(TensorListOptimizer):
         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 | Literal['eig'] = tikhonov
+        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',
@@ -93,21 +100,22 @@ class ScipyMinimize(TensorListOptimizer):
             'newton-cg', 'dogleg', 'trust-ncg', 'trust-krylov', 'trust-exact'
         ]
 
+        # jac in scipy is '2-point', '3-point', 'cs', True or None.
         if self.jac == 'autograd':
             if self.use_jac_autograd: self.jac = True
             else: self.jac = None
 
 
-    def _hess(self, x: np.ndarray, params: TensorList, closure: _ClosureType): # type:ignore
+    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 = hessian([value], wrt = params) # type:ignore
-        Hmat =  hessian_list_to_mat(H)
-        regularize_hessian_(Hmat, self.tikhonov)
-        return Hmat.detach().cpu().numpy()
+            _, 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()
 
-    def _objective(self, x: np.ndarray, params: TensorList, closure: _ClosureType):
+    def _objective(self, x: np.ndarray, params: TensorList, closure):
         # set params to x
         params.from_vec_(torch.from_numpy(x).to(params[0], copy=False))
 
@@ -118,7 +126,7 @@ class ScipyMinimize(TensorListOptimizer):
         return _ensure_float(closure(False))
 
     @torch.no_grad
-    def step(self, closure: _ClosureType): # type:ignore # pylint:disable = signature-differs
+    def step(self, closure: Closure):# pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
         params = self.get_params()
 
         # determine hess argument
@@ -130,7 +138,7 @@ class ScipyMinimize(TensorListOptimizer):
         x0 = params.to_vec().detach().cpu().numpy()
 
         # make bounds
-        lb, ub = self.get_group_keys('lb', 'ub', cls=list)
+        lb, ub = self.group_vals('lb', 'ub', cls=list)
         bounds = []
         for p, l, u in zip(params, lb, ub):
             bounds.extend([(l, u)] * p.numel())
@@ -156,8 +164,8 @@ class ScipyMinimize(TensorListOptimizer):
 
 
 
-class ScipyRoot(TensorListOptimizer):
-    """Find a root of a vector function (UNTESTED!).
+class ScipyRootOptimization(Optimizer):
+    """Optimization via using scipy.root on gradients, mainly for experimenting!
 
     Args:
         params: iterable of parameters to optimize or dicts defining parameter groups.
@@ -196,94 +204,11 @@ class ScipyRoot(TensorListOptimizer):
         self.jac = jac
         if self.jac == 'autograd': self.jac = True
 
-    def _objective(self, x: np.ndarray, params: TensorList, closure: _ClosureType):
-        # set params to x
-        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
-
-        # return value and maybe gradients
-        if self.jac:
-            with torch.enable_grad():
-                value = closure(False)
-                if not isinstance(value, torch.Tensor):
-                    raise TypeError(f"Autograd jacobian requires closure to return torch.Tensor, got {type(value)}")
-            jac = jacobian_list_to_vec(jacobian([value], wrt=params))
-            return _ensure_numpy(value), jac.detach().cpu().numpy()
-        return _ensure_numpy(closure(False))
-
-    @torch.no_grad
-    def step(self, closure: _ClosureType): # type:ignore # pylint:disable = signature-differs
-        params = self.get_params()
-
-        x0 = params.to_vec().detach().cpu().numpy()
-
-        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.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
-        return res.fun
-
-
-class ScipyRootOptimization(TensorListOptimizer):
-    """Optimization via finding roots of the gradient with `scipy.optimize.root` (for experiments, won't work well on most problems).
-
-    Args:
-        params: iterable of parameters to optimize or dicts defining parameter groups.
-        method (str, optional): one of methods from https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.root.html#scipy.optimize.root. Defaults to 'hybr'.
-        tol (float | None, optional): tolerance. Defaults to None.
-        callback (_type_, optional): callback. Defaults to None.
-        options (_type_, optional): options for optimizer. Defaults to None.
-        jac (Literal['2, optional): jacobian calculation method. Defaults to 'autograd'.
-        tikhonov (float | Literal['eig'], optional): tikhonov regularization (only for 'hybr' and 'lm'). Defaults to 0.
-        add_loss (float, optional): adds loss value to jacobian multiplied by this to try to avoid finding maxima. Defaults to 0.
-        mul_loss (float, optional): multiplies jacobian by loss value multiplied by this to try to avoid finding maxima. Defaults to 0.
-    """
-    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',
-        tikhonov: float | Literal['eig'] = 0,
-        add_loss: float = 0,
-        mul_loss: float = 0,
-    ):
-        super().__init__(params, {})
-        self.method = method
-        self.tol = tol
-        self.callback = callback
-        self.options = options
-        self.value = None
-        self.tikhonov: float | Literal['eig'] = tikhonov
-        self.add_loss = add_loss
-        self.mul_loss = mul_loss
-
-        self.jac = jac == 'autograd'
-
         # 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: TensorList, closure: _ClosureType):
+    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))
 
@@ -293,23 +218,16 @@ class ScipyRootOptimization(TensorListOptimizer):
                 self.value = closure(False)
                 if not isinstance(self.value, torch.Tensor):
                     raise TypeError(f"Autograd jacobian requires closure to return torch.Tensor, got {type(self.value)}")
-                jac_list, hess_list = jacobian_and_hessian([self.value], wrt=params)
-            jac = jacobian_list_to_vec(jac_list)
-            hess = hessian_list_to_mat(hess_list)
-            regularize_hessian_(hess, self.tikhonov)
-            if self.mul_loss != 0: jac *= self.value * self.mul_loss
-            if self.add_loss != 0: jac += self.value * self.add_loss
-            return jac.detach().cpu().numpy(), hess.detach().cpu().numpy()
+                g, H = jacobian_and_hessian_mat_wrt([self.value], wrt=params)
+            return g.detach().cpu().numpy(), H.detach().cpu().numpy()
 
         # return the gradients
         with torch.enable_grad(): self.value = closure()
         jac = params.ensure_grad_().grad.to_vec()
-        if self.mul_loss != 0: jac *= self.value * self.mul_loss
-        if self.add_loss != 0: jac += self.value * self.add_loss
         return jac.detach().cpu().numpy()
 
     @torch.no_grad
-    def step(self, closure: _ClosureType): # type:ignore # pylint:disable = signature-differs
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
         params = self.get_params()
 
         x0 = params.to_vec().detach().cpu().numpy()
@@ -325,9 +243,11 @@ class ScipyRootOptimization(TensorListOptimizer):
         )
 
         params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
-        return self.value
+        return res.fun
+
 
-class ScipyDE(TensorListOptimizer):
+
+class ScipyDE(Optimizer):
     """Use scipy.minimize.differential_evolution as pytorch optimizer. Note that this performs full minimization on each step,
     so usually you would want to perform a single step. This also requires bounds to be specified.
 
@@ -374,12 +294,12 @@ class ScipyDE(TensorListOptimizer):
         self._kwargs = kwargs
         self._lb, self._ub = bounds
 
-    def _objective(self, x: np.ndarray, params: TensorList, closure: _ClosureType):
+    def _objective(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))
         return _ensure_float(closure(False))
 
     @torch.no_grad
-    def step(self, closure: _ClosureType): # type:ignore # pylint:disable = signature-differs
+    def step(self, closure: Closure): # pylint:disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
         params = self.get_params()
 
         x0 = params.to_vec().detach().cpu().numpy()
@@ -396,44 +316,47 @@ class ScipyDE(TensorListOptimizer):
         return res.fun
 
 
-class ScipyMinimizeSubspace(Modular):
-    """for experiments and won't work well on most problems.
 
-    explanation - optimizes in a small subspace using scipy.optimize.minimize, but doesnt seem to work well"""
+class ScipyDualAnnealing(Optimizer):
     def __init__(
         self,
         params,
-        projections: Projection | abc.Iterable[Projection] = (
-            Proj2Masks(5),
-            ProjNormalize(
-                ProjGrad(),
-            )
-        ),
-        method=None,
-        lb = None,
-        ub = None,
-        constraints=(),
-        tol=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 = '2-point',
+        bounds: tuple[float, float],
+        maxiter=1000,
+        minimizer_kwargs=None,
+        initial_temp=5230.0,
+        restart_temp_ratio=2.0e-5,
+        visit=2.62,
+        accept=-5.0,
+        maxfun=1e7,
+        rng=None,
+        no_local_search=False,
     ):
+        super().__init__(params, {})
 
-        scopt = WrapClosure(
-                ScipyMinimize,
-                method = method,
-                lb = lb,
-                ub = ub,
-                constraints = constraints,
-                tol = tol,
-                callback = callback,
-                options = options,
-                jac = jac,
-                hess = hess
-            )
-        modules = [
-            Subspace(scopt, projections),
-        ]
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['bounds'], kwargs['__class__']
+        self._kwargs = kwargs
+        self._lb, self._ub = bounds
+
+    def _objective(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))
+        return _ensure_float(closure(False))
+
+    @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()
+        bounds = [(self._lb, self._ub)] * len(x0)
+
+        res = scipy.optimize.dual_annealing(
+            partial(self._objective, params = params, closure = closure),
+            x0 = x0,
+            bounds=bounds,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return res.fun
 
-        super().__init__(params, modules)

torchzero/utils/__init__.py

@@ -0,0 +1,27 @@
+from . import tensorlist as tl
+from .compile import _optional_compiler, benchmark_compile_cpu, benchmark_compile_cuda, set_compilation, enable_compilation
+from .numberlist import NumberList
+from .optimizer import (
+    Init,
+    ListLike,
+    Optimizer,
+    ParamFilter,
+    get_group_vals,
+    get_params,
+    get_state_vals,
+    grad_at_params,
+    grad_vec_at_params,
+    loss_at_params,
+    loss_grad_at_params,
+    loss_grad_vec_at_params,
+)
+from .params import (
+    Params,
+    _add_defaults_to_param_groups_,
+    _add_updates_grads_to_param_groups_,
+    _copy_param_groups,
+    _make_param_groups,
+)
+from .python_tools import flatten, generic_eq, reduce_dim
+from .tensorlist import TensorList, as_tensorlist, Distributions, generic_clamp, generic_numel, generic_vector_norm, generic_zeros_like, generic_randn_like
+from .torch_tools import tofloat, tolist, tonumpy, totensor, vec_to_tensors, vec_to_tensors_, set_storage_