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

torchzero/optim/wrappers/nlopt.py

@@ -69,7 +69,7 @@ def _ensure_tensor(x):
 inf = float('inf')
 Closure = Callable[[bool], Any]
 
-class NLOptOptimizer(Optimizer):
+class NLOptWrapper(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
@@ -96,9 +96,9 @@ class NLOptOptimizer(Optimizer):
         self,
         params,
         algorithm: int | _ALGOS_LITERAL,
-        maxeval: int | None,
         lb: float | None = None,
         ub: float | None = None,
+        maxeval: int | None = 10000, # None can stall on some algos and because they are threaded C you can't even interrupt them
         stopval: float | None = None,
         ftol_rel: float | None = None,
         ftol_abs: float | None = None,
@@ -122,22 +122,33 @@ class NLOptOptimizer(Optimizer):
         self._last_loss = None
 
     def _f(self, x: np.ndarray, grad: np.ndarray, closure, params: TensorList):
-        t = _ensure_tensor(x)
-        if t is None:
+        if self.raised:
             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)
-            grad[:] = params.ensure_grad_().grad.to_vec().reshape(grad.shape).detach().cpu().numpy()
+            return np.inf
+        try:
+            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)
+                grad[:] = params.ensure_grad_().grad.to_vec().reshape(grad.shape).detach().cpu().numpy()
+                return self._last_loss
+
+            self._last_loss = _ensure_float(closure(False))
             return self._last_loss
-
-        self._last_loss = _ensure_float(closure(False))
-        return self._last_loss
+        except Exception as e:
+            self.e = e
+            self.raised = True
+            if self.opt is not None: self.opt.force_stop()
+            return np.inf
 
     @torch.no_grad
     def step(self, closure: Closure): # pylint: disable = signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
+        self.e = None
+        self.raised = False
         params = self.get_params()
 
         # make bounds
@@ -175,6 +186,9 @@ class NLOptOptimizer(Optimizer):
         except Exception as e:
             raise e from None
 
+        if x is not None: params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        if self.e is not None: raise self.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/optuna.py

@@ -0,0 +1,70 @@
+import typing
+from collections import abc
+
+import numpy as np
+import torch
+
+import optuna
+
+from ...utils import Optimizer
+
+def silence_optuna():
+    optuna.logging.set_verbosity(optuna.logging.WARNING)
+
+def _ensure_float(x) -> float:
+    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
+    if isinstance(x, np.ndarray): return float(x.item())
+    return float(x)
+
+
+class OptunaSampler(Optimizer):
+    """Optimize your next SOTA model using hyperparameter optimization.
+
+    Note - optuna is surprisingly scalable to large number of parameters (up to 10,000), despite literally requiring a for-loop because it only supports scalars. Default TPESampler is good for BBO. Maybe not for NNs...
+
+    Args:
+        params: iterable of parameters to optimize or dicts defining parameter groups.
+        lb (float): lower bounds.
+        ub (float): upper bounds.
+        sampler (optuna.samplers.BaseSampler | type[optuna.samplers.BaseSampler] | None, optional): sampler. Defaults to None.
+        silence (bool, optional): makes optuna not write a lot of very useful information to console. Defaults to True.
+    """
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        sampler: "optuna.samplers.BaseSampler | type[optuna.samplers.BaseSampler] | None" = None,
+        silence: bool = True,
+    ):
+        if silence: silence_optuna()
+        super().__init__(params, lb=lb, ub=ub)
+
+        if isinstance(sampler, type): sampler = sampler()
+        self.sampler = sampler
+        self.study = None
+
+    @torch.no_grad
+    def step(self, closure):
+
+        params = self.get_params()
+        if self.study is None:
+            self.study = optuna.create_study(sampler=self.sampler)
+
+        # some optuna samplers use torch
+        with torch.enable_grad():
+            trial = self.study.ask()
+
+            suggested = []
+            for gi,g in enumerate(self.param_groups):
+                for pi,p in enumerate(g['params']):
+                    lb, ub =  g['lb'], g['ub']
+                    suggested.extend(trial.suggest_float(f'g{gi}_p{pi}_w{i}', lb, ub) for i in range(p.numel()))
+
+        vec = torch.as_tensor(suggested).to(params[0])
+        params.from_vec_(vec)
+
+        loss = closure()
+        with torch.enable_grad(): self.study.tell(trial, loss)
+
+        return loss

torchzero/optim/wrappers/scipy.py

@@ -11,9 +11,9 @@ 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):
+def _ensure_float(x) -> float:
     if isinstance(x, torch.Tensor): return x.detach().cpu().item()
-    if isinstance(x, np.ndarray): return x.item()
+    if isinstance(x, np.ndarray): return float(x.item())
     return float(x)
 
 def _ensure_numpy(x):
@@ -139,9 +139,11 @@ class ScipyMinimize(Optimizer):
 
         # make bounds
         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())
+        bounds = None
+        if any(b is not None for b in lb) or any(b is not None for b in ub):
+            bounds = []
+            for p, l, u in zip(params, lb, ub):
+                bounds.extend([(l, u)] * p.numel())
 
         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
@@ -265,7 +267,8 @@ class ScipyDE(Optimizer):
     def __init__(
         self,
         params,
-        bounds: tuple[float,float],
+        lb: float,
+        ub: float,
         strategy: Literal['best1bin', 'best1exp', 'rand1bin', 'rand1exp', 'rand2bin', 'rand2exp',
             'randtobest1bin', 'randtobest1exp', 'currenttobest1bin', 'currenttobest1exp',
             'best2exp', 'best2bin'] = 'best1bin',
@@ -287,12 +290,11 @@ class ScipyDE(Optimizer):
         integrality = None,
 
     ):
-        super().__init__(params, {})
+        super().__init__(params, lb=lb, ub=ub)
 
         kwargs = locals().copy()
-        del kwargs['self'], kwargs['params'], kwargs['bounds'], kwargs['__class__']
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], 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))
@@ -303,7 +305,11 @@ class ScipyDE(Optimizer):
         params = self.get_params()
 
         x0 = params.to_vec().detach().cpu().numpy()
-        bounds = [(self._lb, self._ub)] * len(x0)
+
+        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())
 
         res = scipy.optimize.differential_evolution(
             partial(self._objective, params = params, closure = closure),
@@ -321,7 +327,8 @@ class ScipyDualAnnealing(Optimizer):
     def __init__(
         self,
         params,
-        bounds: tuple[float, float],
+        lb: float,
+        ub: float,
         maxiter=1000,
         minimizer_kwargs=None,
         initial_temp=5230.0,
@@ -332,23 +339,25 @@ class ScipyDualAnnealing(Optimizer):
         rng=None,
         no_local_search=False,
     ):
-        super().__init__(params, {})
+        super().__init__(params, lb=lb, ub=ub)
 
         kwargs = locals().copy()
-        del kwargs['self'], kwargs['params'], kwargs['bounds'], kwargs['__class__']
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], 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]
+    def step(self, closure: Closure):
         params = self.get_params()
 
         x0 = params.to_vec().detach().cpu().numpy()
-        bounds = [(self._lb, self._ub)] * len(x0)
+        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())
 
         res = scipy.optimize.dual_annealing(
             partial(self._objective, params = params, closure = closure),
@@ -360,3 +369,145 @@ class ScipyDualAnnealing(Optimizer):
         params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
         return res.fun
 
+
+
+class ScipySHGO(Optimizer):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        constraints = None,
+        n: int = 100,
+        iters: int = 1,
+        callback = None,
+        minimizer_kwargs = None,
+        options = None,
+        sampling_method: str = 'simplicial',
+    ):
+        super().__init__(params, lb=lb, ub=ub)
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    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):
+        params = self.get_params()
+
+        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())
+
+        res = scipy.optimize.shgo(
+            partial(self._objective, params = params, closure = closure),
+            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
+
+
+class ScipyDIRECT(Optimizer):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        maxfun: int | None = 1000,
+        maxiter: int = 1000,
+        eps: float = 0.0001,
+        locally_biased: bool = True,
+        f_min: float = -np.inf,
+        f_min_rtol: float = 0.0001,
+        vol_tol: float = 1e-16,
+        len_tol: float = 0.000001,
+        callback = None,
+    ):
+        super().__init__(params, lb=lb, ub=ub)
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    def _objective(self, x: np.ndarray, params: TensorList, closure) -> float:
+        if self.raised: return np.inf
+        try:
+            params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+            return _ensure_float(closure(False))
+        except Exception as e:
+            # he he he ha, I found a way to make exceptions work in fcmaes and scipy direct
+            self.e = e
+            self.raised = True
+            return np.inf
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        self.raised = False
+        self.e = None
+
+        params = self.get_params()
+
+        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())
+
+        res = scipy.optimize.direct(
+            partial(self._objective, params=params, closure=closure),
+            bounds=bounds,
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+
+        if self.e is not None: raise self.e from None
+        return res.fun
+
+
+
+
+class ScipyBrute(Optimizer):
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        Ns: int = 20,
+        full_output: int = 0,
+        finish = scipy.optimize.fmin,
+        disp: bool = False,
+        workers: int = 1
+    ):
+        super().__init__(params, lb=lb, ub=ub)
+
+        kwargs = locals().copy()
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        self._kwargs = kwargs
+
+    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):
+        params = self.get_params()
+
+        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())
+
+        x0 = scipy.optimize.brute(
+            partial(self._objective, params = params, closure = closure),
+            ranges=bounds,
+            **self._kwargs
+        )
+        params.from_vec_(torch.from_numpy(x0).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return None

torchzero/utils/__init__.py

@@ -9,11 +9,7 @@ from .optimizer import (
     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,
+    unpack_states,
 )
 from .params import (
     Params,
@@ -22,6 +18,6 @@ from .params import (
     _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 .python_tools import flatten, generic_eq, generic_ne, reduce_dim, unpack_dicts
+from .tensorlist import TensorList, as_tensorlist, Distributions, generic_clamp, generic_numel, generic_vector_norm, generic_zeros_like, generic_randn_like, generic_finfo_eps
 from .torch_tools import tofloat, tolist, tonumpy, totensor, vec_to_tensors, vec_to_tensors_, set_storage_

torchzero/utils/derivatives.py

@@ -2,6 +2,7 @@ from collections.abc import Iterable, Sequence
 
 import torch
 import torch.autograd.forward_ad as fwAD
+from typing import Literal
 
 from .torch_tools import swap_tensors_no_use_count_check, vec_to_tensors
 
@@ -157,7 +158,7 @@ def hessian_mat(
     method="func",
     vectorize=False,
     outer_jacobian_strategy="reverse-mode",
-):
+) -> torch.Tensor:
     """
     returns hessian matrix for parameters (as if they were flattened and concatenated into a vector).
 
@@ -189,7 +190,7 @@ def hessian_mat(
         return loss
 
     if method == 'func':
-        return torch.func.hessian(func)(torch.cat([p.view(-1) for p in params]).detach().requires_grad_(create_graph))
+        return torch.func.hessian(func)(torch.cat([p.view(-1) for p in params]).detach().requires_grad_(create_graph)) # pyright:ignore[reportReturnType]
 
     if method == 'autograd.functional':
         return torch.autograd.functional.hessian(
@@ -198,7 +199,7 @@ def hessian_mat(
             create_graph=create_graph,
             vectorize=vectorize,
             outer_jacobian_strategy=outer_jacobian_strategy,
-        )
+        ) # pyright:ignore[reportReturnType]
     raise ValueError(method)
 
 def jvp(fn, params: Iterable[torch.Tensor], tangent: Iterable[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
@@ -510,4 +511,4 @@ def hvp_fd_forward(
     torch._foreach_div_(hvp_, h)
 
     if normalize: torch._foreach_mul_(hvp_, vec_norm)
-    return loss, hvp_
+    return loss, hvp_

torchzero/utils/linalg/__init__.py

@@ -2,4 +2,4 @@ from .matrix_funcs import inv_sqrt_2x2, eigvals_func, singular_vals_func, matrix
 from .orthogonalize import gram_schmidt
 from .qr import qr_householder
 from .svd import randomized_svd
-from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve
+from .solve import cg, nystrom_approximation, nystrom_sketch_and_solve, steihaug_toint_cg