torchzero 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

torchzero/optim/wrappers/fcmaes.py

@@ -0,0 +1,97 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import fcmaes
+import fcmaes.optimizer
+import fcmaes.retry
+import numpy as np
+import torch
+
+from ...utils import Optimizer, TensorList
+
+Closure = Callable[[bool], Any]
+
+
+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)
+
+def silence_fcmaes():
+    fcmaes.retry.logger.disable('fcmaes')
+
+class FcmaesWrapper(Optimizer):
+    """Use fcmaes as pytorch optimizer. Particularly fcmaes has BITEOPT which appears to win in many benchmarks.
+
+    Note that this performs full minimization on each step, so only perform one step with this.
+
+    Args:
+        params (_type_): _description_
+        lb (float): _description_
+        ub (float): _description_
+        optimizer (fcmaes.optimizer.Optimizer | None, optional): _description_. Defaults to None.
+        max_evaluations (int | None, optional): _description_. Defaults to 50000.
+        value_limit (float | None, optional): _description_. Defaults to np.inf.
+        num_retries (int | None, optional): _description_. Defaults to 1.
+        workers (int, optional): _description_. Defaults to 1.
+        popsize (int | None, optional): _description_. Defaults to 31.
+        capacity (int | None, optional): _description_. Defaults to 500.
+        stop_fitness (float | None, optional): _description_. Defaults to -np.inf.
+        statistic_num (int | None, optional): _description_. Defaults to 0.
+    """
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        optimizer: fcmaes.optimizer.Optimizer | None = None,
+        max_evaluations: int | None = 50000,
+        value_limit: float | None = np.inf,
+        num_retries: int | None = 1,
+        workers: int = 1,
+        popsize: int | None = 31,
+        capacity: int | None = 500,
+        stop_fitness: float | None = -np.inf,
+        statistic_num: int | None = 0
+    ):
+        super().__init__(params, lb=lb, ub=ub)
+        silence_fcmaes()
+        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:
+            # ha 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 = fcmaes.retry.minimize(
+            partial(self._objective, params=params, closure=closure), # pyright:ignore[reportArgumentType]
+            bounds=bounds, # pyright:ignore[reportArgumentType]
+            **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
+

torchzero/optim/wrappers/mads.py

@@ -0,0 +1,90 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import torch
+from mads.mads import orthomads
+
+from ...utils import Optimizer, TensorList
+
+
+def _ensure_float(x):
+    if isinstance(x, torch.Tensor): return x.detach().cpu().item()
+    if isinstance(x, np.ndarray): return x.item()
+    return float(x)
+
+def _ensure_numpy(x):
+    if isinstance(x, torch.Tensor): return x.detach().cpu()
+    if isinstance(x, np.ndarray): return x
+    return np.array(x)
+
+
+Closure = Callable[[bool], Any]
+
+
+class MADS(Optimizer):
+    """Use mads.orthomads 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.
+
+    Args:
+        params (params): params
+        lb (float): lower bounds
+        ub (float): upper bounds
+        dp (float, optional): Initial poll size as percent of bounds. Defaults to 0.1.
+        dm (float, optional): Initial mesh size as percent of bounds. Defaults to 0.01.
+        dp_tol (_type_, optional): Minimum poll size stopping criteria. Defaults to -float('inf').
+        nitermax (_type_, optional): Maximum objective function evaluations. Defaults to float('inf').
+        displog (bool, optional): whether to show log. Defaults to False.
+        savelog (bool, optional): whether to save log. Defaults to False.
+
+    """
+    def __init__(
+        self,
+        params,
+        lb: float,
+        ub: float,
+        dp = 0.1,
+        dm = 0.01,
+        dp_tol = -float('inf'),
+        nitermax = float('inf'),
+        displog = False,
+        savelog = False,
+    ):
+        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()
+
+        x0 = params.to_vec().detach().cpu().numpy()
+
+        lb, ub = self.group_vals('lb', 'ub', cls=list)
+        bounds_lower = []
+        bounds_upper = []
+        for p, l, u in zip(params, lb, ub):
+            bounds_lower.extend([l] * p.numel())
+            bounds_upper.extend([u] * p.numel())
+
+        f, x = orthomads(
+            design_variables=x0,
+            bounds_upper=np.asarray(bounds_upper),
+            bounds_lower=np.asarray(bounds_lower),
+            objective_function=partial(self._objective, params = params, closure = closure),
+            **self._kwargs
+        )
+
+        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        return f
+

torchzero/optim/wrappers/nevergrad.py

@@ -9,12 +9,12 @@ import nevergrad as ng
 from ...utils import Optimizer
 
 
-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)
 
-class NevergradOptimizer(Optimizer):
+class NevergradWrapper(Optimizer):
     """Use nevergrad optimizer as pytorch optimizer.
     Note that it is recommended to specify `budget` to the number of iterations you expect to run,
     as some nevergrad optimizers will error without it.
@@ -56,7 +56,7 @@ class NevergradOptimizer(Optimizer):
         self.budget = budget
 
     @torch.no_grad
-    def step(self, closure): # type:ignore # pylint:disable=signature-differs
+    def step(self, closure): # pylint:disable=signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
         params = self.get_params()
         if self.opt is None:
             ng_params = []

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 (_type_): parameters
+        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):
@@ -265,7 +265,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 +288,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 +303,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 +325,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 +337,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 +367,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 .python_tools import flatten, generic_eq, reduce_dim, unpack_dicts
 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_

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
 
@@ -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_