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

torchzero/modules/line_search/line_search.py

@@ -8,7 +8,7 @@ from typing import Any, Literal
 import numpy as np
 import torch
 
-from ...core import Module, Target, Var
+from ...core import Module,  Objective
 from ...utils import tofloat, set_storage_
 from ..functional import clip_by_finfo
 
@@ -139,7 +139,7 @@ class LineSearchBase(Module, ABC):
         for c, n in zip(params, new_params):
             set_storage_(c, n)
 
-    def _loss(self, step_size: float, var: Var, closure, params: list[torch.Tensor],
+    def _loss(self, step_size: float, var: Objective, closure, params: list[torch.Tensor],
               update: list[torch.Tensor], backward:bool=False) -> float:
 
         # if step_size is 0, we might already know the loss
@@ -165,16 +165,16 @@ class LineSearchBase(Module, ABC):
         # if evaluated loss at step size 0, set it to var.loss
         if step_size == 0:
             var.loss = loss
-            if backward: var.grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
+            if backward: var.grads = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
 
         return tofloat(loss)
 
-    def _loss_derivative_gradient(self, step_size: float, var: Var, closure,
+    def _loss_derivative_gradient(self, step_size: float, var: Objective, closure,
                          params: list[torch.Tensor], update: list[torch.Tensor]):
         # if step_size is 0, we might already know the derivative
-        if (var.grad is not None) and (step_size == 0):
+        if (var.grads is not None) and (step_size == 0):
             loss = self._loss(step_size=step_size,var=var,closure=closure,params=params,update=update,backward=False)
-            derivative = - sum(t.sum() for t in torch._foreach_mul(var.grad, update))
+            derivative = - sum(t.sum() for t in torch._foreach_mul(var.grads, update))
 
         else:
             # loss with a backward pass sets params.grad
@@ -184,79 +184,79 @@ class LineSearchBase(Module, ABC):
             derivative = - sum(t.sum() for t in torch._foreach_mul([p.grad if p.grad is not None
                                                                     else torch.zeros_like(p) for p in params], update))
 
-        assert var.grad is not None
-        return loss, tofloat(derivative), var.grad
+        assert var.grads is not None
+        return loss, tofloat(derivative), var.grads
 
-    def _loss_derivative(self, step_size: float, var: Var, closure,
+    def _loss_derivative(self, step_size: float, var: Objective, closure,
                          params: list[torch.Tensor], update: list[torch.Tensor]):
         return self._loss_derivative_gradient(step_size=step_size, var=var,closure=closure,params=params,update=update)[:2]
 
-    def evaluate_f(self, step_size: float, var: Var, backward:bool=False):
+    def evaluate_f(self, step_size: float, var: Objective, backward:bool=False):
         """evaluate function value at alpha `step_size`."""
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update(),backward=backward)
+        return self._loss(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_updates(),backward=backward)
 
-    def evaluate_f_d(self, step_size: float, var: Var):
+    def evaluate_f_d(self, step_size: float, var: Objective):
         """evaluate function value and directional derivative in the direction of the update at step size `step_size`."""
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss_derivative(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
+        return self._loss_derivative(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_updates())
 
-    def evaluate_f_d_g(self, step_size: float, var: Var):
+    def evaluate_f_d_g(self, step_size: float, var: Objective):
         """evaluate function value, directional derivative, and gradient list at step size `step_size`."""
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return self._loss_derivative_gradient(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_update())
+        return self._loss_derivative_gradient(step_size=step_size, var=var, closure=closure, params=var.params,update=var.get_updates())
 
-    def make_objective(self, var: Var, backward:bool=False):
+    def make_objective(self, var: Objective, backward:bool=False):
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss, var=var, closure=closure, params=var.params, update=var.get_update(), backward=backward)
+        return partial(self._loss, var=var, closure=closure, params=var.params, update=var.get_updates(), backward=backward)
 
-    def make_objective_with_derivative(self, var: Var):
+    def make_objective_with_derivative(self, var: Objective):
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss_derivative, var=var, closure=closure, params=var.params, update=var.get_update())
+        return partial(self._loss_derivative, var=var, closure=closure, params=var.params, update=var.get_updates())
 
-    def make_objective_with_derivative_and_gradient(self, var: Var):
+    def make_objective_with_derivative_and_gradient(self, var: Objective):
         closure = var.closure
         if closure is None: raise RuntimeError('line search requires closure')
-        return partial(self._loss_derivative_gradient, var=var, closure=closure, params=var.params, update=var.get_update())
+        return partial(self._loss_derivative_gradient, var=var, closure=closure, params=var.params, update=var.get_updates())
 
     @abstractmethod
-    def search(self, update: list[torch.Tensor], var: Var) -> float:
+    def search(self, update: list[torch.Tensor], var: Objective) -> float:
         """Finds the step size to use"""
 
     @torch.no_grad
-    def step(self, var: Var) -> Var:
+    def apply(self, objective: Objective) -> Objective:
         self._reset()
 
-        params = var.params
+        params = objective.params
         self._initial_params = [p.clone() for p in params]
-        update = var.get_update()
+        update = objective.get_updates()
 
         try:
-            step_size = self.search(update=update, var=var)
+            step_size = self.search(update=update, var=objective)
         except MaxLineSearchItersReached:
             step_size = self._best_step_size
 
         step_size = clip_by_finfo(step_size, torch.finfo(update[0].dtype))
 
         # set loss_approx
-        if var.loss_approx is None: var.loss_approx = self._lowest_loss
+        if objective.loss_approx is None: objective.loss_approx = self._lowest_loss
 
         # if this is last module, directly update parameters to avoid redundant operations
-        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.set_step_size_(step_size, params=params, update=update)
 
-            var.stop = True; var.skip_update = True
-            return var
+            objective.stop = True; objective.skip_update = True
+            return objective
 
         # revert parameters and multiply update by step size
         self.set_step_size_(0, params=params, update=update)
-        torch._foreach_mul_(var.update, step_size)
-        return var
+        torch._foreach_mul_(objective.updates, step_size)
+        return objective
 
 
 

torchzero/modules/line_search/strong_wolfe.py

@@ -284,8 +284,8 @@ class StrongWolfe(LineSearchBase):
             'init_value', 'init', 'c1', 'c2', 'a_max', 'maxiter', 'maxzoom',
             'maxeval', 'interpolation', 'adaptive', 'plus_minus', 'fallback', 'tol_change')(self.defaults)
 
-        dir = as_tensorlist(var.get_update())
-        grad_list = var.get_grad()
+        dir = as_tensorlist(var.get_updates())
+        grad_list = var.get_grads()
 
         g_0 = -sum(t.sum() for t in torch._foreach_mul(grad_list, dir))
         f_0 = var.get_loss(False)

torchzero/modules/misc/debug.py

@@ -11,9 +11,9 @@ class PrintUpdate(Module):
         defaults = dict(text=text, print_fn=print_fn)
         super().__init__(defaults)
 
-    def step(self, var):
-        self.defaults["print_fn"](f'{self.defaults["text"]}{var.update}')
-        return var
+    def apply(self, objective):
+        self.defaults["print_fn"](f'{self.defaults["text"]}{objective.updates}')
+        return objective
 
 class PrintShape(Module):
     """Prints shapes of the update."""
@@ -21,10 +21,10 @@ class PrintShape(Module):
         defaults = dict(text=text, print_fn=print_fn)
         super().__init__(defaults)
 
-    def step(self, var):
-        shapes = [u.shape for u in var.update] if var.update is not None else None
+    def apply(self, objective):
+        shapes = [u.shape for u in objective.updates] if objective.updates is not None else None
         self.defaults["print_fn"](f'{self.defaults["text"]}{shapes}')
-        return var
+        return objective
 
 class PrintParams(Module):
     """Prints current update."""
@@ -32,9 +32,9 @@ class PrintParams(Module):
         defaults = dict(text=text, print_fn=print_fn)
         super().__init__(defaults)
 
-    def step(self, var):
-        self.defaults["print_fn"](f'{self.defaults["text"]}{var.params}')
-        return var
+    def apply(self, objective):
+        self.defaults["print_fn"](f'{self.defaults["text"]}{objective.params}')
+        return objective
 
 
 class PrintLoss(Module):
@@ -43,6 +43,6 @@ class PrintLoss(Module):
         defaults = dict(text=text, print_fn=print_fn)
         super().__init__(defaults)
 
-    def step(self, var):
-        self.defaults["print_fn"](f'{self.defaults["text"]}{var.get_loss(False)}')
-        return var
+    def apply(self, objective):
+        self.defaults["print_fn"](f'{self.defaults["text"]}{objective.get_loss(False)}')
+        return objective

torchzero/modules/misc/escape.py

@@ -3,7 +3,7 @@ import math
 from typing import Literal
 import torch
 
-from ...core import Modular, Module, Var, Chainable
+from ...core import Modular, Module, Objective, Chainable
 from ...utils import NumberList, TensorList
 
 
@@ -15,11 +15,11 @@ class EscapeAnnealing(Module):
 
 
     @torch.no_grad
-    def step(self, var):
-        closure = var.closure
+    def apply(self, objective):
+        closure = objective.closure
         if closure is None: raise RuntimeError("Escape requries closure")
 
-        params = TensorList(var.params)
+        params = TensorList(objective.params)
         settings = self.settings[params[0]]
         max_region = self.get_settings(params, 'max_region', cls=NumberList)
         max_iter = settings['max_iter']
@@ -41,7 +41,7 @@ class EscapeAnnealing(Module):
         self.global_state['n_bad'] = n_bad
 
         # no progress
-        f_0 = var.get_loss(False)
+        f_0 = objective.get_loss(False)
         if n_bad >= n_tol:
             for i in range(1, max_iter+1):
                 alpha = max_region * (i / max_iter)
@@ -51,12 +51,12 @@ class EscapeAnnealing(Module):
                 f_star = closure(False)
 
                 if math.isfinite(f_star) and f_star < f_0-1e-12:
-                    var.update = None
-                    var.stop = True
-                    var.skip_update = True
-                    return var
+                    objective.updates = None
+                    objective.stop = True
+                    objective.skip_update = True
+                    return objective
 
                 params.sub_(pert)
 
             self.global_state['n_bad'] = 0
-        return var
+        return objective

torchzero/modules/misc/gradient_accumulation.py

@@ -3,74 +3,6 @@ import torch
 from ...core import Chainable, Module
 
 
-# class GradientAccumulation(Module):
-#     """Uses :code:`n` steps to accumulate gradients, after :code:`n` gradients have been accumulated, they are passed to :code:`modules` and parameters are updates.
-
-#     Accumulating gradients for :code:`n` steps is equivalent to increasing batch size by :code:`n`. Increasing the batch size
-#     is more computationally efficient, but sometimes it is not feasible due to memory constraints.
-
-#     .. note::
-#         Technically this can accumulate any inputs, including updates generated by previous modules. As long as this module is first, it will accumulate the gradients.
-
-#     Args:
-#         modules (Chainable): modules that perform a step every :code:`n` steps using the accumulated gradients.
-#         n (int): number of gradients to accumulate.
-#         mean (bool, optional): if True, uses mean of accumulated gradients, otherwise uses sum. Defaults to True.
-#         stop (bool, optional):
-#             this module prevents next modules from stepping unless :code:`n` gradients have been accumulate. Setting this argument to False disables that. Defaults to True.
-
-#     Examples:
-#         Adam with gradients accumulated for 16 batches.
-
-#         .. code-block:: python
-
-#             opt = tz.Modular(
-#                 model.parameters(),
-#                 tz.m.GradientAccumulation(
-#                     [tz.m.Adam(), tz.m.LR(1e-2)],
-#                     n=16
-#                 )
-#             )
-
-#     """
-#     def __init__(self, modules: Chainable, n: int, mean=True, stop=True):
-#         defaults = dict(n=n, mean=mean, stop=stop)
-#         super().__init__(defaults)
-#         self.set_child('modules', modules)
-
-
-#     @torch.no_grad
-#     def step(self, var):
-#         accumulator = self.get_state(var.params, 'accumulator')
-#         settings = self.defaults
-#         n = settings['n']; mean = settings['mean']; stop = settings['stop']
-#         step = self.global_state['step'] = self.global_state.get('step', 0) + 1
-
-#         # add update to accumulator
-#         torch._foreach_add_(accumulator, var.get_update())
-
-#         # step with accumulated updates
-#         if step % n == 0:
-#             if mean:
-#                 torch._foreach_div_(accumulator, n)
-
-#             var.update = [a.clone() for a in accumulator]
-#             var = self.children['modules'].step(var)
-
-#             # zero accumulator
-#             torch._foreach_zero_(accumulator)
-
-#         else:
-#             # prevent update
-#             if stop:
-#                 var.update = None
-#                 var.stop=True
-#                 var.skip_update=True
-
-#         return var
-
-
-
 
 class GradientAccumulation(Module):
     """Uses ``n`` steps to accumulate gradients, after ``n`` gradients have been accumulated, they are passed to :code:`modules` and parameters are updates.
@@ -106,21 +38,21 @@ class GradientAccumulation(Module):
 
 
     @torch.no_grad
-    def step(self, var):
-        accumulator = self.get_state(var.params, 'accumulator')
+    def apply(self, objective):
+        accumulator = self.get_state(objective.params, 'accumulator')
         settings = self.defaults
         n = settings['n']; mean = settings['mean']; stop = settings['stop']
-        step = self.global_state['step'] = self.global_state.get('step', 0) + 1
+        step = self.increment_counter("step", 0)
 
         # add update to accumulator
-        torch._foreach_add_(accumulator, var.get_update())
+        torch._foreach_add_(accumulator, objective.get_updates())
 
         # step with accumulated updates
-        if step % n == 0:
+        if (step + 1) % n == 0:
             if mean:
                 torch._foreach_div_(accumulator, n)
 
-            var.update = accumulator
+            objective.updates = accumulator
 
             # zero accumulator
             self.clear_state_keys('accumulator')
@@ -128,9 +60,9 @@ class GradientAccumulation(Module):
         else:
             # prevent update
             if stop:
-                var.update = None
-                var.stop=True
-                var.skip_update=True
+                objective.updates = None
+                objective.stop=True
+                objective.skip_update=True
 
-        return var
+        return objective
 

torchzero/modules/misc/homotopy.py

@@ -13,27 +13,27 @@ class HomotopyBase(Module):
         """transform the loss"""
 
     @torch.no_grad
-    def step(self, var):
-        if var.loss is not None:
-            var.loss = self.loss_transform(var.loss)
+    def apply(self, objective):
+        if objective.loss is not None:
+            objective.loss = self.loss_transform(objective.loss)
 
-        closure = var.closure
+        closure = objective.closure
         if closure is None: raise RuntimeError("SquareHomotopy requires closure")
 
         def homotopy_closure(backward=True):
             if backward:
                 with torch.enable_grad():
                     loss = self.loss_transform(closure(False))
-                    grad = torch.autograd.grad(loss, var.params, allow_unused=True)
-                    for p,g in zip(var.params, grad):
+                    grad = torch.autograd.grad(loss, objective.params, allow_unused=True)
+                    for p,g in zip(objective.params, grad):
                         p.grad = g
             else:
                 loss = self.loss_transform(closure(False))
 
             return loss
 
-        var.closure = homotopy_closure
-        return var
+        objective.closure = homotopy_closure
+        return objective
 
 class SquareHomotopy(HomotopyBase):
     def __init__(self): super().__init__()
@@ -57,3 +57,11 @@ class LambdaHomotopy(HomotopyBase):
         super().__init__(defaults)
 
     def loss_transform(self, loss): return self.defaults['fn'](loss)
+
+class FixedLossHomotopy(HomotopyBase):
+    def __init__(self, value: float = 1):
+        defaults = dict(value=value)
+        super().__init__(defaults)
+
+    def loss_transform(self, loss): return loss / loss.detach().clip(min=torch.finfo(loss.dtype).tiny * 2)
+