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

torchzero/modules/misc/multistep.py

@@ -2,49 +2,49 @@ from collections.abc import Iterable
 
 import torch
 
-from ...core import Chainable, Module, Var
+from ...core import Chainable, Module, Objective
 from ...utils import TensorList
 
-def _sequential_step(self: Module, var: Var, sequential: bool):
-    params = var.params
+def _sequential_step(self: Module, objective: Objective, sequential: bool):
+    params = objective.params
     steps = self.settings[params[0]]['steps']
 
-    if sequential: modules = self.get_children_sequence() * steps
+    if sequential: modules: list[Module] = self.get_children_sequence() * steps
     else: modules = [self.children['module']] * steps
 
-    if var.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
+    if objective.closure is None and len(modules) > 1: raise ValueError('Multistep and Sequential require closure')
 
     # store original params unless this is last module and can update params directly
     params_before_steps = [p.clone() for p in params]
 
     # first step - pass var as usual
-    var = modules[0].step(var)
-    new_var = var
+    objective = modules[0].step(objective)
+    new_objective = objective
 
     # subsequent steps - update parameters and create new var
     if len(modules) > 1:
         for m in modules[1:]:
 
             # update params
-            if (not new_var.skip_update):
+            if (not new_objective.skip_update):
                 # if new_var.last_module_lrs is not None:
                 #     torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
-                torch._foreach_sub_(params, new_var.get_update())
+                torch._foreach_sub_(params, new_objective.get_updates())
 
             # create new var since we are at a new point, that means grad, update and loss will be None
-            new_var = Var(params=new_var.params, closure=new_var.closure,
-                            model=new_var.model, current_step=new_var.current_step + 1)
+            new_objective = Objective(params=new_objective.params, closure=new_objective.closure,
+                            model=new_objective.model, current_step=new_objective.current_step + 1)
 
             # step
-            new_var = m.step(new_var)
+            new_objective = m.step(new_objective)
 
         # final parameter update
-        if (not new_var.skip_update):
+        if (not new_objective.skip_update):
             # if new_var.last_module_lrs is not None:
             #     torch._foreach_mul_(new_var.get_update(), new_var.last_module_lrs)
 
-            torch._foreach_sub_(params, new_var.get_update())
+            torch._foreach_sub_(params, new_objective.get_updates())
 
     # if last module, update is applied so return new var
     # if params_before_steps is None:
@@ -53,13 +53,13 @@ def _sequential_step(self: Module, var: Var, sequential: bool):
     #     return new_var
 
     # otherwise use parameter difference as update
-    var.update = list(torch._foreach_sub(params_before_steps, params))
+    objective.updates = list(torch._foreach_sub(params_before_steps, params))
     for p, bef in zip(params, params_before_steps):
         p.set_(bef) # pyright:ignore[reportArgumentType]
-    return var
+    return objective
 
 class Multistep(Module):
-    """Performs :code:`steps` inner steps with :code:`module` per each step.
+    """Performs ``steps`` inner steps with ``module`` per each step.
 
     The update is taken to be the parameter difference between parameters before and after the inner loop."""
     def __init__(self, module: Chainable, steps: int):
@@ -68,11 +68,11 @@ class Multistep(Module):
         self.set_child('module', module)
 
     @torch.no_grad
-    def step(self, var):
-        return _sequential_step(self, var, sequential=False)
+    def apply(self, objective):
+        return _sequential_step(self, objective, sequential=False)
 
 class Sequential(Module):
-    """On each step, this sequentially steps with :code:`modules` :code:`steps` times.
+    """On each step, this sequentially steps with ``modules`` ``steps`` times.
 
     The update is taken to be the parameter difference between parameters before and after the inner loop."""
     def __init__(self, modules: Iterable[Chainable], steps: int=1):
@@ -81,28 +81,28 @@ class Sequential(Module):
         self.set_children_sequence(modules)
 
     @torch.no_grad
-    def step(self, var):
-        return _sequential_step(self, var, sequential=True)
+    def apply(self, objective):
+        return _sequential_step(self, objective, sequential=True)
 
 
 class NegateOnLossIncrease(Module):
-    """Uses an extra forward pass to evaluate loss at :code:`parameters+update`,
-    if loss is larger than at :code:`parameters`,
-    the update is set to 0 if :code:`backtrack=False` and to :code:`-update` otherwise"""
+    """Uses an extra forward pass to evaluate loss at ``parameters+update``,
+    if loss is larger than at ``parameters``,
+    the update is set to 0 if ``backtrack=False`` and to ``-update`` otherwise"""
     def __init__(self, backtrack=False):
         defaults = dict(backtrack=backtrack)
         super().__init__(defaults=defaults)
 
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def apply(self, objective):
+        closure = objective.closure
         if closure is None: raise RuntimeError('NegateOnLossIncrease requires closure')
         backtrack = self.defaults['backtrack']
 
-        update = var.get_update()
-        f_0 = var.get_loss(backward=False)
+        update = objective.get_updates()
+        f_0 = objective.get_loss(backward=False)
 
-        torch._foreach_sub_(var.params, update)
+        torch._foreach_sub_(objective.params, update)
         f_1 = closure(False)
 
         if f_1 <= f_0:
@@ -111,15 +111,15 @@ class NegateOnLossIncrease(Module):
             #     var.skip_update = True
             #     return var
 
-            torch._foreach_add_(var.params, update)
-            return var
+            torch._foreach_add_(objective.params, update)
+            return objective
 
-        torch._foreach_add_(var.params, update)
+        torch._foreach_add_(objective.params, update)
         if backtrack:
-            torch._foreach_neg_(var.update)
+            torch._foreach_neg_(objective.updates)
         else:
-            torch._foreach_zero_(var.update)
-        return var
+            torch._foreach_zero_(objective.updates)
+        return objective
 
 
 class Online(Module):
@@ -129,7 +129,7 @@ class Online(Module):
 
     Online L-BFGS with Backtracking line search
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.Online(tz.m.LBFGS()),
         tz.m.Backtracking()
@@ -138,57 +138,56 @@ class Online(Module):
 
     Online L-BFGS trust region
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.TrustCG(tz.m.Online(tz.m.LBFGS()))
     )
     ```
 
     """
-    def __init__(self, *modules: Module,):
+    def __init__(self, module: Module,):
         super().__init__()
-
-        self.set_child('module', modules)
+        self.set_child('module', module)
 
     @torch.no_grad
-    def update(self, var):
-        closure = var.closure
+    def update(self, objective):
+        closure = objective.closure
         if closure is None: raise ValueError("Closure must be passed for Online")
 
         step = self.global_state.get('step', 0) + 1
         self.global_state['step'] = step
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         p_cur = params.clone()
         p_prev = self.get_state(params, 'p_prev', cls=TensorList)
 
         module = self.children['module']
-        var_c = var.clone(clone_update=False)
+        var_c = objective.clone(clone_updates=False)
 
         # on 1st step just step and store previous params
         if step == 1:
             p_prev.copy_(params)
 
             module.update(var_c)
-            var.update_attrs_from_clone_(var_c)
+            objective.update_attrs_from_clone_(var_c)
             return
 
         # restore previous params and update
-        var_prev = Var(params=params, closure=closure, model=var.model, current_step=var.current_step)
+        prev_objective = Objective(params=params, closure=closure, model=objective.model, current_step=objective.current_step)
         params.set_(p_prev)
         module.reset_for_online()
-        module.update(var_prev)
+        module.update(prev_objective)
 
         # restore current params and update
         params.set_(p_cur)
         p_prev.copy_(params)
         module.update(var_c)
-        var.update_attrs_from_clone_(var_c)
+        objective.update_attrs_from_clone_(var_c)
 
     @torch.no_grad
-    def apply(self, var):
+    def apply(self, objective):
         module = self.children['module']
-        return module.apply(var.clone(clone_update=False))
+        return module.apply(objective.clone(clone_updates=False))
 
-    def get_H(self, var):
-        return self.children['module'].get_H(var)
+    def get_H(self, objective):
+        return self.children['module'].get_H(objective)

torchzero/modules/misc/regularization.py

@@ -1,14 +1,14 @@
 import torch
 
-from ...core import Chainable, Module, Target, Transform
+from ...core import Chainable, Module,  Transform
 from ...core.reformulation import Reformulation
-from ...utils import Distributions, NumberList, TensorList
+from ...utils import Distributions, Metrics, NumberList, TensorList, evaluate_metric
 
 
 class Dropout(Transform):
     """Applies dropout to the update.
 
-    For each weight the update to that weight has :code:`p` probability to be set to 0.
+    For each weight the update to that weight has ``p`` probability to be set to 0.
     This can be used to implement gradient dropout or update dropout depending on placement.
 
     Args:
@@ -18,36 +18,37 @@ class Dropout(Transform):
         target (Target, optional): what to set on var, refer to documentation. Defaults to 'update'.
 
 
-    Examples:
-        Gradient dropout.
+    ### Examples:
 
-        .. code-block:: python
+    Gradient dropout.
 
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Dropout(0.5),
-                tz.m.Adam(),
-                tz.m.LR(1e-3)
-            )
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.Dropout(0.5),
+        tz.m.Adam(),
+        tz.m.LR(1e-3)
+    )
+    ```
 
-        Update dropout.
+    Update dropout.
 
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Adam(),
-                tz.m.Dropout(0.5),
-                tz.m.LR(1e-3)
-            )
+    ``python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.Adam(),
+        tz.m.Dropout(0.5),
+        tz.m.LR(1e-3)
+    )
+    ```
 
     """
-    def __init__(self, p: float = 0.5, graft: bool=False, target: Target = 'update'):
+    def __init__(self, p: float = 0.5, graft: bool=False):
         defaults = dict(p=p, graft=graft)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults)
 
     @torch.no_grad
-    def apply_tensors(self, tensors, params, grads, loss, states, settings):
+    def multi_tensor_apply(self, tensors, params, grads, loss, states, settings):
         tensors = TensorList(tensors)
         p = NumberList(s['p'] for s in settings)
         graft = settings[0]['graft']
@@ -67,32 +68,31 @@ class WeightDropout(Module):
     """
     Changes the closure so that it evaluates loss and gradients with random weights replaced with 0.
 
-    Dropout can be disabled for a parameter by setting :code:`use_dropout=False` in corresponding parameter group.
+    Dropout can be disabled for a parameter by setting ``use_dropout=False`` in corresponding parameter group.
 
     Args:
         p (float, optional): probability that any weight is replaced with 0. Defaults to 0.5.
-        graft (bool, optional):
-            if True, parameters after dropout are rescaled to have the same norm as before dropout. Defaults to False.
     """
-    def __init__(self, p: float = 0.5, graft: bool = True):
-        defaults = dict(p=p, graft=graft, use_dropout=True)
+    def __init__(self, p: float = 0.5):
+        defaults = dict(p=p, use_dropout=True)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def update(self, objective):
+        closure = objective.closure
         if closure is None: raise RuntimeError('WeightDropout requires closure')
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         p = NumberList(self.settings[p]['p'] for p in params)
 
         # create masks
         mask = []
-        for p, m in zip(params, mask):
+        for p in params:
             prob = self.settings[p]['p']
             use_dropout = self.settings[p]['use_dropout']
             if use_dropout: mask.append(_bernoulli_like(p, prob))
             else: mask.append(torch.ones_like(p))
 
+        # create a closure that evaluates masked parameters
         @torch.no_grad
         def dropout_closure(backward=True):
             orig_params = params.clone()
@@ -104,15 +104,14 @@ class WeightDropout(Module):
             params.copy_(orig_params)
             return loss
 
-        var.closure = dropout_closure
-        return var
+        objective.closure = dropout_closure
 
 
 class PerturbWeights(Module):
     """
     Changes the closure so that it evaluates loss and gradients at weights perturbed by a random perturbation.
 
-    Can be disabled for a parameter by setting :code:`perturb=False` in corresponding parameter group.
+    Can be disabled for a parameter by setting ``perturb=False`` in corresponding parameter group.
 
     Args:
         alpha (float, optional): multiplier for perturbation magnitude. Defaults to 0.1.
@@ -120,15 +119,22 @@ class PerturbWeights(Module):
         distribution (bool, optional):
             distribution of the random perturbation. Defaults to False.
     """
-    def __init__(self, alpha: float = 0.1, relative:bool=True, distribution:Distributions = 'normal'):
-        defaults = dict(alpha=alpha, relative=relative, distribution=distribution, perturb=True)
+
+    def __init__(
+        self,
+        alpha: float = 0.1,
+        relative: bool = True,
+        distribution: Distributions = "normal",
+        metric: Metrics = "mad",
+    ):
+        defaults = dict(alpha=alpha, relative=relative, distribution=distribution, metric=metric, perturb=True)
         super().__init__(defaults)
 
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def update(self, objective):
+        closure = objective.closure
         if closure is None: raise RuntimeError('WeightDropout requires closure')
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
 
         # create perturbations
         perts = []
@@ -140,7 +146,7 @@ class PerturbWeights(Module):
 
             alpha = settings['alpha']
             if settings['relative']:
-                alpha *= p.abs().mean()
+                alpha *= evaluate_metric(p, settings["metric"])
 
             distribution = self.settings[p]['distribution'].lower()
             if distribution in ('normal', 'gaussian'):
@@ -163,5 +169,4 @@ class PerturbWeights(Module):
             params.sub_(perts)
             return loss
 
-        var.closure = perturbed_closure
-        return var
+        objective.closure = perturbed_closure

torchzero/modules/misc/split.py

@@ -1,54 +1,53 @@
-import warnings
 from collections.abc import Callable, Sequence, Iterable
 from typing import cast
 
 import torch
 
-from ...core import Chainable, Module, Var
+from ...core import Chainable, Module, Objective
 
 
 def _split(
     module: Module,
     idxs,
     params,
-    var: Var,
+    objective: Objective,
 ):
     split_params = [p for i,p in enumerate(params) if i in idxs]
 
     split_grad = None
-    if var.grad is not None:
-        split_grad = [g for i,g in enumerate(var.grad) if i in idxs]
+    if objective.grads is not None:
+        split_grad = [g for i,g in enumerate(objective.grads) if i in idxs]
 
     split_update = None
-    if var.update is not None:
-        split_update = [u for i,u in enumerate(var.update) if i in idxs]
+    if objective.updates is not None:
+        split_update = [u for i,u in enumerate(objective.updates) if i in idxs]
 
-    split_var = var.clone(clone_update=False, parent=var)
-    split_var.params = split_params
-    split_var.grad = split_grad
-    split_var.update = split_update
+    split_obj = objective.clone(clone_updates=False, parent=objective)
+    split_obj.params = split_params
+    split_obj.grads = split_grad
+    split_obj.updates = split_update
 
-    split_var = module.step(split_var)
+    split_obj = module.step(split_obj)
 
     # those should be set due to var being parent
-    if split_var.grad is not None:
-        assert var.grad is not None
+    if split_obj.grads is not None:
+        assert objective.grads is not None
 
-    if split_var.loss is not None:
-        assert var.loss is not None
+    if split_obj.loss is not None:
+        assert objective.loss is not None
 
-    if split_var.update is not None:
+    if split_obj.updates is not None:
 
         # make sure update is set, it will be filled with ``true`` and ``false`` tensors
-        if var.update is None:
-            if var.grad is None: var.update = [cast(torch.Tensor, None) for _ in var.params]
-            else: var.update = [g.clone() for g in var.grad]
+        if objective.updates is None:
+            if objective.grads is None: objective.updates = [cast(torch.Tensor, None) for _ in objective.params]
+            else: objective.updates = [g.clone() for g in objective.grads]
 
         # set all tensors from this split
-        for idx, u in zip(idxs, split_var.update):
-            var.update[idx] = u
+        for idx, u in zip(idxs, split_obj.updates):
+            objective.updates[idx] = u
 
-    return var
+    return objective
 
 _SingleFilter = Callable[[torch.Tensor], bool] | torch.Tensor | Iterable[torch.Tensor] | torch.nn.Module | Iterable[torch.nn.Module]
 Filter = _SingleFilter | Iterable[_SingleFilter]
@@ -82,7 +81,7 @@ class Split(Module):
     Muon with Adam fallback using same hyperparams as https://github.com/KellerJordan/Muon
 
     ```python
-    opt = tz.Modular(
+    opt = tz.Optimizer(
         model.parameters(),
         tz.m.NAG(0.95),
         tz.m.Split(
@@ -101,9 +100,12 @@ class Split(Module):
         if true is not None: self.set_child('true', true)
         if false is not None: self.set_child('false', false)
 
-    def step(self, var):
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
 
-        params = var.params
+    def step(self, objective):
+
+        params = objective.params
         filter = _make_filter(self.settings[params[0]]['filter'])
 
         true_idxs = []
@@ -114,10 +116,10 @@ class Split(Module):
 
         if 'true' in self.children and len(true_idxs) > 0:
             true = self.children['true']
-            var = _split(true, idxs=true_idxs, params=params, var=var)
+            objective = _split(true, idxs=true_idxs, params=params, objective=objective)
 
         if 'false' in self.children and len(false_idxs) > 0:
             false = self.children['false']
-            var = _split(false, idxs=false_idxs, params=params, var=var)
+            objective = _split(false, idxs=false_idxs, params=params, objective=objective)
 
-        return var
+        return objective

torchzero/modules/misc/switch.py

@@ -14,20 +14,21 @@ class Alternate(Module):
     Args:
         steps (int | Iterable[int], optional): number of steps to perform with each module. Defaults to 1.
 
-    Examples:
-        Alternate between Adam, SignSGD and RMSprop
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Alternate(
-                    tz.m.Adam(),
-                    [tz.m.SignSGD(), tz.m.Mul(0.5)],
-                    tz.m.RMSprop(),
-                ),
-                tz.m.LR(1e-3),
-            )
+    ### Examples:
+    Alternate between Adam, SignSGD and RMSprop
+
+    ```python
+
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.Alternate(
+            tz.m.Adam(),
+            [tz.m.SignSGD(), tz.m.Mul(0.5)],
+            tz.m.RMSprop(),
+        ),
+        tz.m.LR(1e-3),
+    )
+    ```
     """
     LOOP = True
     def __init__(self, *modules: Chainable, steps: int | Iterable[int] = 1):
@@ -43,14 +44,17 @@ class Alternate(Module):
         self.global_state['current_module_idx'] = 0
         self.global_state['steps_to_next'] = steps[0] if isinstance(steps, list) else steps
 
+    def update(self, objective): raise RuntimeError
+    def apply(self, objective): raise RuntimeError
+
     @torch.no_grad
-    def step(self, var):
+    def step(self, objective):
         # get current module
         current_module_idx = self.global_state.setdefault('current_module_idx', 0)
         module = self.children[f'module_{current_module_idx}']
 
         # step
-        var = module.step(var.clone(clone_update=False))
+        objective = module.step(objective.clone(clone_updates=False))
 
         # number of steps until next module
         steps = self.defaults['steps']
@@ -72,28 +76,29 @@ class Alternate(Module):
 
             self.global_state['steps_to_next'] = steps[self.global_state['current_module_idx']]
 
-        return var
+        return objective
 
 class Switch(Alternate):
-    """After :code:`steps` steps switches to the next module.
+    """After ``steps`` steps switches to the next module.
 
     Args:
         steps (int | Iterable[int]): Number of steps to perform with each module.
 
-    Examples:
-        Start with Adam, switch to L-BFGS after 1000th step and Truncated Newton on 2000th step.
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Switch(
-                    [tz.m.Adam(), tz.m.LR(1e-3)],
-                    [tz.m.LBFGS(), tz.m.Backtracking()],
-                    [tz.m.NewtonCG(maxiter=20), tz.m.Backtracking()],
-                    steps = (1000, 2000)
-                )
-            )
+    ### Examples:
+
+    Start with Adam, switch to L-BFGS after 1000th step and Truncated Newton on 2000th step.
+
+    ```python
+    opt = tz.Optimizer(
+        model.parameters(),
+        tz.m.Switch(
+            [tz.m.Adam(), tz.m.LR(1e-3)],
+            [tz.m.LBFGS(), tz.m.Backtracking()],
+            [tz.m.NewtonCG(maxiter=20), tz.m.Backtracking()],
+            steps = (1000, 2000)
+        )
+    )
+    ```
     """
 
     LOOP = False