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

torchzero/core/transform.py

@@ -25,7 +25,8 @@ class Transform(Module, ABC):
     def __init__(
         self,
         defaults: dict[str,Any] | None,
-        uses_grad: bool,
+        uses_grad: bool = False,
+        uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
         scale_first: bool = False,
@@ -35,49 +36,48 @@ class Transform(Module, ABC):
         super().__init__(defaults)
         self._target: Target = target
         self._uses_grad = uses_grad
+        self._uses_loss = uses_loss
         self._concat_params = concat_params
         self._update_freq = update_freq
         self._scale_first = scale_first
         self._inner = inner
 
-    def update(
+    def update_tensors(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]],
     ) -> None:
-        """Updates this transform. By default does nothing - if logic is in `apply` method."""
+        """update function, this shouldn't be called directly. Updates this module."""
 
     @abstractmethod
-    def apply(
+    def apply_tensors(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]],
     ) -> Sequence[torch.Tensor]:
-        """Applies the update rule to `tensors`."""
+        """apply function, this shouldn't be called directly. Applies the update rule to `tensors` and returns them.
+        If possible, this shouldn't modify the internal state of this transform."""
 
     @final
     @torch.no_grad
-    def transform(
+    def transform_update(
         self,
         tensors: list[torch.Tensor],
         params: list[torch.Tensor],
         grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         states: list[dict[str, Any]],
         settings: Sequence[Mapping[str, Any]] | None,
-    ) -> list[torch.Tensor]:
-        """Applies this transform to an arbitrary sequence of tensors."""
-        un_tensors = tensors
-        un_params = params
-        un_grads = grads
+    ) -> None:
+        """Updates this transform from an arbitrary sequence of tensors."""
         if self._concat_params:
             tensors = [torch.cat([t.ravel() for t in tensors])]
             params = [torch.cat([p.ravel() for p in params])]
@@ -86,24 +86,61 @@ class Transform(Module, ABC):
         if settings is None:
             settings = [self.defaults for _ in tensors]
 
-        step = self.global_state.get('__step', 0)
+        step = self.global_state.get('__step', 0) # that way it gets reset correctly
+        self.global_state['__step'] = step + 1
+
         num = len(tensors)
         states = states[:num]
         settings = settings[:num]
 
-        update_freq = self._update_freq
-        scale_first = self._scale_first
         scale_factor = 1
 
         # scaling factor for 1st step
-        if scale_first and step == 0:
+        if self._scale_first and step == 0:
             # initial step size guess from pytorch LBFGS
             scale_factor = 1 / TensorList(tensors).abs().global_sum().clip(min=1)
             scale_factor = scale_factor.clip(min=torch.finfo(tensors[0].dtype).eps)
 
         # update transform
-        if step % update_freq == 0:
-            self.update(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+        if step % self._update_freq == 0:
+            self.update_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+        # store for transform_apply
+        self.global_state["__tensors"] = tensors
+        self.global_state["__params"] = params
+        self.global_state["__grads"] = grads
+        self.global_state["__scale_factor"] = scale_factor
+
+
+    @final
+    @torch.no_grad
+    def transform_apply(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+        states: list[dict[str, Any]],
+        settings: Sequence[Mapping[str, Any]] | None,
+    ) -> list[torch.Tensor]:
+        """Applies this transform to an arbitrary sequence of tensors.
+        This can be used after ``transform_update`` has been used at least once."""
+
+        if settings is None:
+            settings = [self.defaults for _ in tensors]
+
+        num = len(tensors)
+        states = states[:num]
+        settings = settings[:num]
+
+        un_tensors = tensors
+        un_params = params
+        un_grads = grads
+
+        tensors = self.global_state.pop("__tensors")
+        params  = self.global_state.pop("__params")
+        grads   = self.global_state.pop("__grads")
+        scale_factor = self.global_state.pop("__scale_factor")
 
         # step with inner
         if self._inner is not None:
@@ -112,27 +149,17 @@ class Transform(Module, ABC):
                 tensors = [torch.cat([t.ravel() for t in tensors])]
 
         # apply transform
-        tensors = list(self.apply(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
+        tensors = list(self.apply_tensors(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings))
 
         # scale initial step, when preconditioner might not have been applied
-        if scale_first and step == 0:
+        if self._scale_first and self.global_state['__step'] == 1:
             torch._foreach_mul_(tensors, scale_factor)
 
-        self.global_state['__step'] = step + 1
         if self._concat_params:
             tensors = vec_to_tensors(vec=tensors[0], reference=un_tensors)
         return tensors
 
-
-    @torch.no_grad
-    def keyed_transform(
-        self,
-        tensors: list[torch.Tensor],
-        params: list[torch.Tensor],
-        grads: list[torch.Tensor] | None,
-        loss: torch.Tensor | None,
-    ):
-        """Applies this transform to `tensors`, `params` will be used as keys and need to always point to same tensor objects."""
+    def _get_keyed_states_settings(self, params: list[torch.Tensor]):
         if self._concat_params:
             p = params[0]
             states = [self.state[p]]
@@ -145,41 +172,116 @@ class Transform(Module, ABC):
                 states.append(self.state[p])
                 settings.append(self.settings[p])
 
-        return self.transform(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+        return states, settings
+
+    @final
+    @torch.no_grad
+    def keyed_transform_update(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+    ):
+        """`params` will be used as keys and need to always point to same tensor objects.`"""
+        states, settings = self._get_keyed_states_settings(params)
+        self.transform_update(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+
+    @final
+    @torch.no_grad
+    def keyed_transform_apply(
+        self,
+        tensors: list[torch.Tensor],
+        params: list[torch.Tensor],
+        grads: list[torch.Tensor] | None,
+        loss: torch.Tensor | float | None,
+    ):
+        """`params` will be used as keys and need to always point to same tensor objects.`"""
+        states, settings = self._get_keyed_states_settings(params)
+        return self.transform_apply(tensors=tensors, params=params, grads=grads, loss=loss, states=states, settings=settings)
+
+
+    def pre_step(self, var: Var) -> None:
+        """Logic to run pre-transform, this way transform has access to  Var."""
+    def post_step(self, var: Var) -> None:
+        """Logic to run post-transform, this way transform has access to  Var."""
+
+    def update(self, var: Var):
+        if self._target != 'update':
+            raise ValueError("Target must be 'update' to use `update` and `apply` methods. "
+                             f"With {self._target = } only `step` method can be used.")
+
+        # var may change, therefore current params and grads have to be extracted and passed explicitly
+        update = var.get_update() # this sets loss
+        if self._uses_grad: var.get_grad()
+        if self._uses_loss: var.get_loss(False)
+        params=var.params
+        self.pre_step(var)
+
+        # update
+        self.keyed_transform_update(update, params, var.grad, var.loss)
+
+    def apply(self, var: Var):
+        if self._target != 'update':
+            raise ValueError("Target must be 'update' to use `update` and `apply` methods. "
+                             f"With {self._target = } only `step` method can be used.")
 
-    def step(self, var: Var) -> Var:
         # var may change, therefore current params and grads have to be extracted and passed explicitly
+        update = var.get_update() # this sets loss
         if self._uses_grad: var.get_grad()
+        if self._uses_loss: var.get_loss(False)
+        params=var.params
+
+        # apply
+        var.update = self.keyed_transform_apply(update, params, var.grad, var.loss)
+        self.post_step(var)
+        return var
+
+    def step(self, var: Var) -> Var:
+
+        # var may change, therefore current params and grads have to be extracted and passed explicitly
+        if self._target in ('update', 'update_difference'): var.get_update() # this sets loss
+        if self._uses_grad or self._target == 'grad': var.get_grad()
+        if self._uses_loss: var.get_loss(False)
         params=var.params
+        self.pre_step(var)
 
         # ---------------------------------- update ---------------------------------- #
         if self._target == 'update':
             update = var.get_update()
-            var.update = list(self.keyed_transform(update, params, var.grad, var.loss))
+            self.keyed_transform_update(update, params, var.grad, var.loss)
+            var.update = list(self.keyed_transform_apply(update, params, var.grad, var.loss))
             return var
 
         # ----------------------------------- grad ----------------------------------- #
         if self._target == 'grad':
             grad = var.get_grad()
-            var.grad = list(self.keyed_transform(grad, params, grad, var.loss))
+            self.keyed_transform_update(grad, params, grad, var.loss)
+            var.grad = list(self.keyed_transform_apply(grad, params, grad, var.loss))
             return var
 
         # ------------------------------- params_direct ------------------------------ #
         if self._target == 'params_direct':
-            new_params = self.keyed_transform(var.params, params, var.grad, var.loss)
+            self.keyed_transform_update(var.params, params, var.grad, var.loss)
+            new_params = self.keyed_transform_apply(var.params, params, var.grad, var.loss)
             for p, new_p in zip(var.params, new_params): set_storage_(p, new_p)
             return var
 
         # ----------------------------- params_differnce ----------------------------- #
         if self._target == 'params_difference':
-            new_params = tuple(self.keyed_transform([p.clone() for p in var.params], params, var.grad, var.loss))
+            p_clone = [p.clone() for p in var.params]
+            self.keyed_transform_update(p_clone, params, var.grad, var.loss)
+            new_params = tuple(self.keyed_transform_apply(p_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(var.params, new_params))
             return var
 
         # ----------------------------- update_difference ---------------------------- #
         if self._target == 'update_difference':
             update = var.get_update()
-            new_update = tuple(self.keyed_transform([u.clone() for u in update], params, var.grad, var.loss))
+            u_clone = [u.clone() for u in update]
+            self.keyed_transform_update(u_clone, params, var.grad, var.loss)
+            new_update = tuple(self.keyed_transform_apply(u_clone, params, var.grad, var.loss))
             var.update = list(torch._foreach_sub(update, new_update))
             return var
 
@@ -193,7 +295,8 @@ class Transform(Module, ABC):
                 if backward:
                     loss = original_closure()
                     current_grad = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params]
-                    transformed_grad = list(self.keyed_transform(current_grad, params, var.grad, var.loss))
+                    self.keyed_transform_update(current_grad, params, var.grad, var.loss)
+                    transformed_grad = list(self.keyed_transform_apply(current_grad, params, var.grad, var.loss))
                     for p, g in zip(params, transformed_grad):
                         p.grad = g
 
@@ -203,6 +306,7 @@ class Transform(Module, ABC):
                 return loss
 
             var.closure = transformed_closure
+            self.post_step(var)
             return var
 
         # ---------------------------------- invalid --------------------------------- #
@@ -225,7 +329,8 @@ class TensorwiseTransform(Transform, ABC):
     def __init__(
         self,
         defaults: dict[str,Any] | None,
-        uses_grad: bool,
+        uses_grad: bool = False,
+        uses_loss: bool = False,
         concat_params: bool = False,
         update_freq: int = 1,
         scale_first: bool = False,
@@ -238,6 +343,7 @@ class TensorwiseTransform(Transform, ABC):
             concat_params=concat_params,
             update_freq=update_freq,
             scale_first=scale_first,
+            uses_loss=uses_loss,
             inner=inner,
             target=target,
         )
@@ -247,9 +353,9 @@ class TensorwiseTransform(Transform, ABC):
         tensor: torch.Tensor,
         param: torch.Tensor,
         grad: torch.Tensor | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         state: dict[str, Any],
-        settings: Mapping[str, Any],
+        setting: Mapping[str, Any],
     ) -> None:
         """Updates this transform. By default does nothing - if logic is in `apply` method."""
 
@@ -259,20 +365,20 @@ class TensorwiseTransform(Transform, ABC):
         tensor: torch.Tensor,
         param: torch.Tensor,
         grad: torch.Tensor | None,
-        loss: torch.Tensor | None,
+        loss: torch.Tensor | float | None,
         state: dict[str, Any],
-        settings: Mapping[str, Any],
+        setting: Mapping[str, Any],
     ) -> torch.Tensor:
         """Applies the update rule to `tensor`."""
 
     @final
-    def update(self, tensors, params, grads, loss, states, settings):
+    def update_tensors(self, tensors, params, grads, loss, states, settings):
         if grads is None: grads = [None]*len(tensors)
         for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
             self.update_tensor(t, p, g, loss, state, setting)
 
     @final
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         applied = []
         if grads is None: grads = [None]*len(tensors)
         for t,p,g,state,setting in zip(tensors, params, grads, states, settings):
@@ -284,7 +390,7 @@ def apply_transform(
     tensors: list[torch.Tensor],
     params: list[torch.Tensor],
     grads: list[torch.Tensor] | None,
-    loss: torch.Tensor | None = None,
+    loss: torch.Tensor | float | None = None,
     var: Var | None = None,
     current_step: int = 0,
 ):
@@ -292,9 +398,10 @@ def apply_transform(
         var = Var(params=params, closure=None, model=None, current_step=current_step)
         var.loss = loss
 
-    if isinstance(tfm, Transform):
+    if isinstance(tfm, Transform) and tfm._target == 'update':
         if tfm._uses_grad and grads is None: grads = var.get_grad()
-        return list(tfm.keyed_transform(tensors, params, grads, loss))
+        tfm.keyed_transform_update(tensors, params, grads, loss)
+        return list(tfm.keyed_transform_apply(tensors, params, grads, loss))
 
     if isinstance(tfm, Chain): tfm = tfm.get_children_sequence() # pyright: ignore[reportAssignmentType]
     if isinstance(tfm, Sequence):

torchzero/modules/__init__.py

@@ -1,7 +1,7 @@
 from .clipping import *
 from .grad_approximation import *
 from .line_search import *
-from .lr import *
+from .step_size import *
 from .momentum import *
 from .ops import *
 from .optimizers import *
@@ -11,4 +11,5 @@ from .smoothing import *
 from .weight_decay import *
 from .wrappers import *
 from .second_order import *
-from .higher_order import *
+from .higher_order import *
+from .misc import *

torchzero/modules/clipping/clipping.py

@@ -5,7 +5,7 @@ import math
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList, generic_eq
+from ...utils import NumberList, TensorList
 
 
 def clip_grad_value_(params: Iterable[torch.Tensor], value: float):
@@ -24,7 +24,7 @@ def _clip_norm_(
     min: float | NumberList | None,
     max: float | NumberList | None,
     norm_value: float | NumberList | None,
-    ord: float,
+    ord: float | Literal['mean_abs'],
     dim: int | Sequence[int] | Literal["global"] | None,
     inverse_dims: bool,
     min_size: int,
@@ -54,9 +54,13 @@ def _clip_norm_(
         size = math.prod(tensor.size(d) for d in real_dim)
         if size < min_size: continue
 
-        norm: torch.Tensor = torch.linalg.vector_norm(tensor, ord=ord, dim=real_dim, keepdim=True) # pylint:disable=not-callable
+        if ord == 'mean_abs':
+            norm = tensor.abs().mean(dim=real_dim, keepdim=True)
+        else:
+            norm: torch.Tensor = torch.linalg.vector_norm(tensor, ord=ord, dim=real_dim, keepdim=True) # pylint:disable=not-callable
+
         if norm.numel() == 1 and norm == 0: continue
-        norm = torch.where(norm == 0, 1, norm)
+        norm = torch.where(norm <= 1e-12, 1, norm)
 
         # normalize = True, perform normalization
         norm_v = norm_value[i] if isinstance(norm_value, (list,tuple)) else norm_value
@@ -90,7 +94,7 @@ def _clip_norm_(
 def clip_grad_norm_(
     params: Iterable[torch.Tensor],
     max_norm: float | None,
-    ord: float = 2,
+    ord: float | Literal['mean_abs'] = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 2,
@@ -118,7 +122,7 @@ def clip_grad_norm_(
 def normalize_grads_(
     params: Iterable[torch.Tensor],
     norm_value: float,
-    ord: float = 2,
+    ord: float | Literal['mean_abs'] = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 1,
@@ -145,13 +149,43 @@ def normalize_grads_(
 
 
 class ClipValue(Transform):
-    """Clips update magnitude to be within `(-value, value)` range."""
+    """Clips update magnitude to be within `(-value, value)` range.
+
+    Args:
+        value (float): value to clip to.
+        target (str): refer to :ref:`target argument` in documentation.
+
+    Examples:
+
+        Gradient clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ClipValue(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.ClipValue(1),
+                tz.m.LR(1e-2),
+            )
+
+    """
     def __init__(self, value: float, target: Target = 'update'):
         defaults = dict(value=value)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         value = [s['value'] for s in settings]
         return TensorList(tensors).clip_([-v for v in value], value)
 
@@ -172,21 +206,45 @@ class ClipNorm(Transform):
             minimal numer of elements in a parameter or slice to clip norm. Defaults to 1.
         target (str, optional):
             what this affects.
+
+    Examples:
+
+        Gradient norm clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.ClipNorm(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update norm clipping:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.ClipNorm(1),
+                tz.m.LR(1e-2),
+            )
     """
     def __init__(
         self,
         max_norm: float,
-        ord: float = 2,
+        ord: float | Literal['mean_abs'] = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
         target: Target = "update",
     ):
         defaults = dict(max_norm=max_norm,ord=ord,dim=dim,min_size=min_size,inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         max_norm = NumberList(s['max_norm'] for s in settings)
         ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
         _clip_norm_(
@@ -218,21 +276,45 @@ class Normalize(Transform):
             minimal size of a dimension to normalize along it. Defaults to 1.
         target (str, optional):
             what this affects.
+
+    Examples:
+
+        Gradient normalization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Normalize(1),
+                tz.m.Adam(),
+                tz.m.LR(1e-2),
+            )
+
+        Update normalization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Adam(),
+                tz.m.Normalize(1),
+                tz.m.LR(1e-2),
+            )
     """
     def __init__(
         self,
         norm_value: float = 1,
-        ord: float = 2,
+        ord: float | Literal['mean_abs'] = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
         target: Target = "update",
     ):
         defaults = dict(norm_value=norm_value,ord=ord,dim=dim,min_size=min_size, inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         norm_value = NumberList(s['norm_value'] for s in settings)
         ord, dim, min_size, inverse_dims = itemgetter('ord', 'dim', 'min_size', 'inverse_dims')(settings[0])
 
@@ -299,6 +381,21 @@ class Centralize(Transform):
             if True, the `dims` argument is inverted, and all other dimensions are centralized.
         min_size (int, optional):
             minimal size of a dimension to normalize along it. Defaults to 1.
+
+    Examples:
+
+        Standard gradient centralization:
+
+        .. code-block:: python
+
+            opt = tz.Modular(
+                model.parameters(),
+                tz.m.Centralize(dim=0),
+                tz.m.LR(1e-2),
+            )
+
+    References:
+        - Yong, H., Huang, J., Hua, X., & Zhang, L. (2020). Gradient centralization: A new optimization technique for deep neural networks. In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part I 16 (pp. 635-652). Springer International Publishing. https://arxiv.org/abs/2004.01461
     """
     def __init__(
         self,
@@ -308,10 +405,10 @@ class Centralize(Transform):
         target: Target = "update",
     ):
         defaults = dict(dim=dim,min_size=min_size,inverse_dims=inverse_dims)
-        super().__init__(defaults, uses_grad=False, target=target)
+        super().__init__(defaults, target=target)
 
     @torch.no_grad
-    def apply(self, tensors, params, grads, loss, states, settings):
+    def apply_tensors(self, tensors, params, grads, loss, states, settings):
         dim, min_size, inverse_dims = itemgetter('dim', 'min_size', 'inverse_dims')(settings[0])
 
         _centralize_(tensors_ = TensorList(tensors), dim=dim, inverse_dims=inverse_dims, min_size=min_size)