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

torchzero/modules/clipping/clipping.py

@@ -1,11 +1,13 @@
+import math
+from collections.abc import Iterable, Sequence
 from operator import itemgetter
 from typing import Literal
-from collections.abc import Iterable, Sequence
-import math
+
 import torch
 
 from ...core import Module, Target, Transform
-from ...utils import NumberList, TensorList
+from ...utils import Metrics, NumberList, TensorList
+from ...utils.metrics import _METRICS
 
 
 def clip_grad_value_(params: Iterable[torch.Tensor], value: float):
@@ -24,7 +26,7 @@ def _clip_norm_(
     min: float | NumberList | None,
     max: float | NumberList | None,
     norm_value: float | NumberList | None,
-    ord: float | Literal['mean_abs'],
+    ord: Metrics,
     dim: int | Sequence[int] | Literal["global"] | None,
     inverse_dims: bool,
     min_size: int,
@@ -35,7 +37,7 @@ def _clip_norm_(
             raise ValueError(f'if norm_value is given then min and max must be None got {min = }; {max = }')
 
         # if dim is None: return tensors_.mul_(norm_value / tensors_.norm(ord=ord))
-        if dim == 'global': return tensors_.mul_(norm_value / tensors_.global_vector_norm(ord=ord))
+        if dim == 'global': return tensors_.mul_(norm_value / tensors_.global_metric(ord))
 
     # if dim is None: return tensors_.clip_norm_(min,max,tensorwise=True,ord=ord)
     if dim == 'global': return tensors_.clip_norm_(min,max,tensorwise=False,ord=ord)
@@ -54,8 +56,8 @@ def _clip_norm_(
         size = math.prod(tensor.size(d) for d in real_dim)
         if size < min_size: continue
 
-        if ord == 'mean_abs':
-            norm = tensor.abs().mean(dim=real_dim, keepdim=True)
+        if isinstance(ord, str):
+            norm = _METRICS[ord].evaluate_tensor(tensor, 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
 
@@ -94,7 +96,7 @@ def _clip_norm_(
 def clip_grad_norm_(
     params: Iterable[torch.Tensor],
     max_norm: float | None,
-    ord: float | Literal['mean_abs'] = 2,
+    ord: Metrics = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 2,
@@ -105,7 +107,7 @@ def clip_grad_norm_(
 
     Args:
         params (Iterable[torch.Tensor]): parameters with gradients to clip.
-        value (float): value to clip norm to.
+        max_norm (float): value to clip norm to.
         ord (float, optional): norm order. Defaults to 2.
         dim (int | Sequence[int] | str | None, optional):
             calculates norm along those dimensions.
@@ -122,7 +124,7 @@ def clip_grad_norm_(
 def normalize_grads_(
     params: Iterable[torch.Tensor],
     norm_value: float,
-    ord: float | Literal['mean_abs'] = 2,
+    ord: Metrics = 2,
     dim: int | Sequence[int] | Literal["global"] | None = None,
     inverse_dims: bool = False,
     min_size: int = 1,
@@ -149,35 +151,33 @@ 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.
+        target (str): refer to ``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),
-            )
+    Gradient clipping:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.ClipValue(1),
+        tz.m.Adam(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
+    Update clipping:
+    ```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'):
@@ -193,7 +193,7 @@ class ClipNorm(Transform):
     """Clips update norm to be no larger than `value`.
 
     Args:
-        value (float): value to clip norm to.
+        max_norm (float): value to clip norm to.
         ord (float, optional): norm order. Defaults to 2.
         dim (int | Sequence[int] | str | None, optional):
             calculates norm along those dimensions.
@@ -209,32 +209,30 @@ class ClipNorm(Transform):
 
     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),
-            )
+    Gradient norm clipping:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.ClipNorm(1),
+        tz.m.Adam(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
+    Update norm clipping:
+    ```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 | Literal['mean_abs'] = 2,
+        ord: Metrics = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
@@ -263,7 +261,7 @@ class Normalize(Transform):
     """Normalizes the update.
 
     Args:
-        value (float): desired norm value.
+        norm_value (float): desired norm value.
         ord (float, optional): norm order. Defaults to 2.
         dim (int | Sequence[int] | str | None, optional):
             calculates norm along those dimensions.
@@ -278,33 +276,31 @@ class Normalize(Transform):
             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),
-            )
+    Gradient normalization:
+    ```python
+    opt = tz.Modular(
+        model.parameters(),
+        tz.m.Normalize(1),
+        tz.m.Adam(),
+        tz.m.LR(1e-2),
+    )
+    ```
+
+    Update normalization:
+
+    ```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 | Literal['mean_abs'] = 2,
+        ord: Metrics = 2,
         dim: int | Sequence[int] | Literal["global"] | None = None,
         inverse_dims: bool = False,
         min_size: int = 1,
@@ -370,8 +366,6 @@ class Centralize(Transform):
     """Centralizes the update.
 
     Args:
-        value (float): desired norm value.
-        ord (float, optional): norm order. Defaults to 2.
         dim (int | Sequence[int] | str | None, optional):
             calculates norm along those dimensions.
             If list/tuple, tensors are centralized along all dimensios in `dim` that they have.
@@ -384,18 +378,17 @@ class Centralize(Transform):
 
     Examples:
 
-        Standard gradient centralization:
-
-        .. code-block:: python
-
-            opt = tz.Modular(
-                model.parameters(),
-                tz.m.Centralize(dim=0),
-                tz.m.LR(1e-2),
-            )
+    Standard gradient centralization:
+    ```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
+    - 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,

torchzero/modules/clipping/ema_clipping.py

@@ -5,7 +5,7 @@ from collections.abc import Iterable, Sequence
 import torch
 
 from ...core import Module, Target, Transform, apply_transform, Chainable
-from ...utils import NumberList, TensorList, unpack_dicts, unpack_states
+from ...utils import NumberList, TensorList, unpack_dicts, unpack_states, Metrics
 
 class ClipNormByEMA(Transform):
     """Clips norm to be no larger than the norm of an exponential moving average of past updates.
@@ -25,7 +25,7 @@ class ClipNormByEMA(Transform):
     def __init__(
         self,
         beta=0.99,
-        ord: float = 2,
+        ord: Metrics = 2,
         eps=1e-6,
         tensorwise:bool=True,
         max_ema_growth: float | None = 1.5,
@@ -47,7 +47,7 @@ class ClipNormByEMA(Transform):
         ema.lerp_(tensors, 1-beta)
 
         if tensorwise:
-            ema_norm = ema.norm(ord)
+            ema_norm = ema.metric(ord)
 
             # clip ema norm growth
             if max_ema_growth is not None:
@@ -64,7 +64,7 @@ class ClipNormByEMA(Transform):
             else: denom.clip_(min=1)
 
         else:
-            ema_norm = ema.global_vector_norm(ord)
+            ema_norm = ema.global_metric(ord)
 
             # clip ema norm growth
             if max_ema_growth is not None:
@@ -75,7 +75,7 @@ class ClipNormByEMA(Transform):
                     ema_norm = allowed_norm
                 prev_ema_norm.set_(ema_norm)
 
-            tensors_norm = tensors.global_vector_norm(ord)
+            tensors_norm = tensors.global_metric(ord)
             denom = tensors_norm / ema_norm.clip(min=eps[0])
             if self.NORMALIZE: denom.clip_(min=eps[0])
             else: denom.clip_(min=1)

torchzero/modules/conjugate_gradient/__init__.py

@@ -0,0 +1,11 @@
+from .cg import (
+    DYHS,
+    ConjugateDescent,
+    DaiYuan,
+    FletcherReeves,
+    HagerZhang,
+    HestenesStiefel,
+    LiuStorey,
+    PolakRibiere,
+    ProjectedGradientMethod,
+)