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

torchzero/modules/wrappers/optim_wrapper.py

@@ -3,7 +3,7 @@ from typing import Any
 import torch
 
 from ...core.module import Module
-from ...utils import Params, _copy_param_groups, _make_param_groups
+from ...utils.params import Params, _copy_param_groups, _make_param_groups
 
 
 class Wrap(Module):
@@ -66,8 +66,8 @@ class Wrap(Module):
         return super().set_param_groups(param_groups)
 
     @torch.no_grad
-    def step(self, var):
-        params = var.params
+    def apply(self, objective):
+        params = objective.params
 
         # initialize opt on 1st step
         if self.optimizer is None:
@@ -76,7 +76,7 @@ class Wrap(Module):
             self.optimizer = self._opt_fn(param_groups, *self._opt_args, **self._opt_kwargs)
 
         # set optimizer per-parameter settings
-        if self.defaults["use_param_groups"] and var.modular is not None:
+        if self.defaults["use_param_groups"] and objective.modular is not None:
             for group in self.optimizer.param_groups:
                 first_param = group['params'][0]
                 setting = self.settings[first_param]
@@ -91,19 +91,19 @@ class Wrap(Module):
 
         # set grad to update
         orig_grad = [p.grad for p in params]
-        for p, u in zip(params, var.get_update()):
+        for p, u in zip(params, objective.get_updates()):
             p.grad = u
 
         # if this is last module, simply use optimizer to update parameters
-        if var.modular is not None and self is var.modular.modules[-1]:
+        if objective.modular is not None and self is objective.modular.modules[-1]:
             self.optimizer.step()
 
             # restore grad
             for p, g in zip(params, orig_grad):
                 p.grad = g
 
-            var.stop = True; var.skip_update = True
-            return var
+            objective.stop = True; objective.skip_update = True
+            return objective
 
         # this is not the last module, meaning update is difference in parameters
         # and passed to next module
@@ -111,11 +111,11 @@ class Wrap(Module):
         self.optimizer.step() # step and update params
         for p, g in zip(params, orig_grad):
             p.grad = g
-        var.update = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
+        objective.updates = list(torch._foreach_sub(params_before_step, params)) # set update to difference between params
         for p, o in zip(params, params_before_step):
             p.set_(o) # pyright: ignore[reportArgumentType]
 
-        return var
+        return objective
 
     def reset(self):
         super().reset()

torchzero/modules/zeroth_order/cd.py

@@ -33,13 +33,16 @@ class CD(Module):
         defaults = dict(h=h, grad=grad, adaptive=adaptive, index=index, threepoint=threepoint)
         super().__init__(defaults)
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def step(self, objective):
+        closure = objective.closure
         if closure is None:
             raise RuntimeError("CD requires closure")
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         ndim = params.global_numel()
 
         grad_step_size = self.defaults['grad']
@@ -79,7 +82,7 @@ class CD(Module):
             else:
                 warnings.warn("CD adaptive=True only works with threepoint=True")
 
-        f_0 = var.get_loss(False)
+        f_0 = objective.get_loss(False)
         params.flat_set_lambda_(idx, lambda x: x + h)
         f_p = closure(False)
 
@@ -117,6 +120,6 @@ class CD(Module):
         # ----------------------------- create the update ---------------------------- #
         update = params.zeros_like()
         update.flat_set_(idx, alpha)
-        var.update = update
-        return var
+        objective.updates = update
+        return objective
 

torchzero/optim/root.py

@@ -3,7 +3,7 @@ from collections.abc import Callable
 
 from abc import abstractmethod
 import torch
-from ..modules.higher_order.multipoint import sixth_order_im1, sixth_order_p6, _solve
+from ..modules.second_order.multipoint import sixth_order_3p, sixth_order_5p, two_point_newton, sixth_order_3pm2, _solve
 
 def make_evaluate(f: Callable[[torch.Tensor], torch.Tensor]):
     def evaluate(x, order) -> tuple[torch.Tensor, ...]:
@@ -53,7 +53,7 @@ class Newton(RootBase):
     def one_iteration(self, x, evaluate): return newton(x, evaluate, self.lstsq)
 
 
-class SixthOrderP6(RootBase):
+class SixthOrder3P(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.
@@ -62,4 +62,4 @@ class SixthOrderP6(RootBase):
     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)
+        return sixth_order_3p(x, f, f_j, self.lstsq)

torchzero/optim/utility/split.py

@@ -3,7 +3,8 @@ from collections.abc import Callable, Iterable
 
 import torch
 
-from ...utils import flatten, get_params
+from ...utils import flatten
+from ...utils.optimizer import get_params
 
 class Split(torch.optim.Optimizer):
     """Steps will all `optimizers`, also has a check that they have no duplicate parameters.

torchzero/optim/wrappers/directsearch.py

@@ -7,24 +7,13 @@ import numpy as np
 import torch
 from directsearch.ds import DEFAULT_PARAMS
 
-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)
-
+from ...utils import TensorList
+from .wrapper import WrapperBase
 
 Closure = Callable[[bool], Any]
 
 
-class DirectSearch(Optimizer):
+class DirectSearch(WrapperBase):
     """Use directsearch as pytorch optimizer.
 
     Note that this performs full minimization on each step,
@@ -96,28 +85,23 @@ class DirectSearch(Optimizer):
         del kwargs['self'], kwargs['params'], 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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
 
         res = directsearch.solve(
-            partial(self._objective, params = params, closure = closure),
+            partial(self._f, params=params, closure=closure),
             x0 = x0,
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.f
 
 
 
-class DirectSearchDS(Optimizer):
+class DirectSearchDS(WrapperBase):
     def __init__(
         self,
         params,
@@ -139,26 +123,21 @@ class DirectSearchDS(Optimizer):
         del kwargs['self'], kwargs['params'], 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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
 
         res = directsearch.solve_directsearch(
-            partial(self._objective, params = params, closure = closure),
+            partial(self._f, params = params, closure = closure),
             x0 = x0,
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.f
 
-class DirectSearchProbabilistic(Optimizer):
+class DirectSearchProbabilistic(WrapperBase):
     def __init__(
         self,
         params,
@@ -179,27 +158,22 @@ class DirectSearchProbabilistic(Optimizer):
         del kwargs['self'], kwargs['params'], 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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
 
         res = directsearch.solve_probabilistic_directsearch(
-            partial(self._objective, params = params, closure = closure),
+            partial(self._f, params = params, closure = closure),
             x0 = x0,
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.f
 
 
-class DirectSearchSubspace(Optimizer):
+class DirectSearchSubspace(WrapperBase):
     def __init__(
         self,
         params,
@@ -223,28 +197,23 @@ class DirectSearchSubspace(Optimizer):
         del kwargs['self'], kwargs['params'], 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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
 
         res = directsearch.solve_subspace_directsearch(
-            partial(self._objective, params = params, closure = closure),
+            partial(self._f, params = params, closure = closure),
             x0 = x0,
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.f
 
 
 
-class DirectSearchSTP(Optimizer):
+class DirectSearchSTP(WrapperBase):
     def __init__(
         self,
         params,
@@ -260,21 +229,16 @@ class DirectSearchSTP(Optimizer):
         del kwargs['self'], kwargs['params'], 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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
 
         res = directsearch.solve_stp(
-            partial(self._objective, params = params, closure = closure),
+            partial(self._f, params = params, closure = closure),
             x0 = x0,
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(res.x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.f

torchzero/optim/wrappers/fcmaes.py

@@ -9,20 +9,15 @@ import fcmaes
 import fcmaes.optimizer
 import fcmaes.retry
 
-from ...utils import Optimizer, TensorList
+from ...utils import TensorList
+from .wrapper import WrapperBase
 
 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):
+class FcmaesWrapper(WrapperBase):
     """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.
@@ -42,7 +37,7 @@ class FcmaesWrapper(Optimizer):
             CMA-ES population size used for all CMA-ES runs.
             Not used for differential evolution.
             Ignored if parameter optimizer is defined. Defaults to 31.
-        capacity (int | None, optional): capacity of the evaluation store.. Defaults to 500.
+        capacity (int | None, optional): capacity of the evaluation store. Defaults to 500.
         stop_fitness (float | None, optional):
             Limit for fitness value. optimization runs terminate if this value is reached. Defaults to -np.inf.
         statistic_num (int | None, optional):
@@ -61,46 +56,30 @@ class FcmaesWrapper(Optimizer):
         popsize: int | None = 31,
         capacity: int | None = 500,
         stop_fitness: float | None = -np.inf,
-        statistic_num: int | None = 0
+        statistic_num: int | None = 0,
+        silence: bool = True,
     ):
-        super().__init__(params, lb=lb, ub=ub)
-        silence_fcmaes()
+        super().__init__(params, dict(lb=lb,ub=ub))
+        if silence:
+            silence_fcmaes()
         kwargs = locals().copy()
-        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__']
+        del kwargs['self'], kwargs['params'], kwargs['lb'], kwargs['ub'], kwargs['__class__'], kwargs["silence"]
         self._kwargs = kwargs
         self._kwargs['workers'] = 1
 
-    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())
+        params = TensorList(self._get_params())
+        bounds = self._get_bounds()
 
         res = fcmaes.retry.minimize(
-            partial(self._objective, params=params, closure=closure), # pyright:ignore[reportArgumentType]
+            partial(self._f, 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
+        params.from_vec_(torch.as_tensor(res.x, device = params[0].device, dtype=params[0].dtype))
         return res.fun
 

torchzero/optim/wrappers/mads.py

@@ -6,24 +6,13 @@ 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)
-
+from ...utils import TensorList
+from .wrapper import WrapperBase
 
 Closure = Callable[[bool], Any]
 
 
-class MADS(Optimizer):
+class MADS(WrapperBase):
     """Use mads.orthomads as pytorch optimizer.
 
     Note that this performs full minimization on each step,
@@ -53,37 +42,28 @@ class MADS(Optimizer):
         displog = False,
         savelog = False,
     ):
-        super().__init__(params, lb=lb, ub=ub)
+        super().__init__(params, dict(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()
+        params = TensorList(self._get_params())
+        x0 = params.to_vec().numpy(force=True)
+        lb, ub = self._get_lb_ub()
 
-        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),
+            bounds_upper=np.asarray(ub),
+            bounds_lower=np.asarray(lb),
+            objective_function=partial(self._f, params=params, closure=closure),
             **self._kwargs
         )
 
-        params.from_vec_(torch.from_numpy(x).to(device = params[0].device, dtype=params[0].dtype, copy=False))
+        params.from_vec_(torch.as_tensor(x, device = params[0].device, dtype=params[0].dtype,))
         return f
 

torchzero/optim/wrappers/moors.py

@@ -0,0 +1,66 @@
+from collections.abc import Callable
+from functools import partial
+from typing import Any, Literal
+
+import numpy as np
+import torch
+
+from ...utils import TensorList
+from .wrapper import WrapperBase
+
+Closure = Callable[[bool], Any]
+
+class MoorsWrapper(WrapperBase):
+    """Use moo-rs (pymoors) is PyTorch optimizer.
+
+    Note that this performs full minimization on each step,
+    so usually you would want to perform a single step.
+
+    To use this, define a function that accepts fitness function and number of variables and returns a pymoors algorithm:
+
+    ```python
+    alg_fn = lambda fitness_fn, num_vars: pymoors.Nsga2(
+        fitness_fn=fitness_fn,
+        num_vars=num_vars,
+        num_iterations=100,
+        sampler = pymoors.RandomSamplingFloat(min=-3, max=3),
+        crossover = pymoors.SinglePointBinaryCrossover(),
+        mutation = pymoors.GaussianMutation(gene_mutation_rate=1e-2, sigma=0.1),
+        population_size = 32,
+        num_offsprings = 32,
+    )
+
+    optimizer = MoorsWrapper(model.parameters(), alg_fn)
+    ```
+
+    All algorithms in pymoors have slightly different APIs, refer to their docs.
+
+    """
+    def __init__(
+        self,
+        params,
+        algorithm_fn: Callable[[Callable[[np.ndarray], np.ndarray], int], Any]
+    ):
+        super().__init__(params, {})
+        self._algorithm_fn = algorithm_fn
+
+    def _objective(self, x: np.ndarray, params, closure):
+        fs = []
+        for x_i in x:
+            f_i = self._fs(x_i, params=params, closure=closure)
+            fs.append(f_i)
+        return np.stack(fs, dtype=np.float64) # pymoors needs float64
+
+    @torch.no_grad
+    def step(self, closure: Closure):
+        params = TensorList(self._get_params())
+        objective = partial(self._objective, params=params, closure=closure)
+
+        algorithm = self._algorithm_fn(objective, params.global_numel())
+
+        algorithm.run()
+        pop = algorithm.population
+
+        params.from_vec_(torch.as_tensor(pop.best[0].genes, device = params[0].device, dtype=params[0].dtype,))
+        return pop.best[0].fitness
+

torchzero/optim/wrappers/nevergrad.py

@@ -6,7 +6,7 @@ import torch
 
 import nevergrad as ng
 
-from ...utils import Optimizer
+from .wrapper import WrapperBase
 
 
 def _ensure_float(x) -> float:
@@ -14,7 +14,7 @@ def _ensure_float(x) -> float:
     if isinstance(x, np.ndarray): return float(x.item())
     return float(x)
 
-class NevergradWrapper(Optimizer):
+class NevergradWrapper(WrapperBase):
     """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.
@@ -72,7 +72,7 @@ class NevergradWrapper(Optimizer):
 
     @torch.no_grad
     def step(self, closure): # pylint:disable=signature-differs # pyright:ignore[reportIncompatibleMethodOverride]
-        params = self.get_params()
+        params = self._get_params()
         if self.opt is None:
             ng_params = []
             for group in self.param_groups:
@@ -95,7 +95,7 @@ class NevergradWrapper(Optimizer):
 
         x: ng.p.Tuple = self.opt.ask() # type:ignore
         for cur, new in zip(params, x):
-            cur.set_(torch.from_numpy(new.value).to(dtype=cur.dtype, device=cur.device, copy=False).reshape_as(cur)) # type:ignore
+            cur.set_(torch.as_tensor(new.value, dtype=cur.dtype, device=cur.device).reshape_as(cur)) # type:ignore
 
         loss = closure(False)
         self.opt.tell(x, _ensure_float(loss))