🚸 Provide option to disable debugging (#17)

🚸 Improve type error messages

README.md

@@ -129,3 +129,10 @@ If you need the absolute best performances, the assertions can be disabled with
 ```bash
 python -O your_awesome_code_using_piqa.py
 ```
+
+Alternatively, you can disable PIQA's type assertions within your code with
+
+```python
+from piqa.utils import set_debug
+set_debug(False)
+```

piqa/fsim.py

@@ -21,7 +21,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.color import ColorConv
 from piqa.utils.functional import (
     scharr_kernel,
@@ -308,7 +308,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -339,4 +339,4 @@ def forward(
         # FSIM
         l = fsim(input, target, pc_input, pc_target, kernel=self.kernel, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/gmsd.py

@@ -22,7 +22,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.color import ColorConv
 from piqa.utils.functional import (
     prewitt_kernel,
@@ -219,7 +219,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -239,7 +239,7 @@ def forward(
         # GMSD
         l = gmsd(input, target, kernel=self.kernel, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)
 
 
 class MS_GMSD(nn.Module):
@@ -310,7 +310,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -331,4 +331,4 @@ def forward(
             **self.kwargs,
         )
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/haarpsi.py

@@ -19,7 +19,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.color import ColorConv
 from piqa.utils.functional import (
     haar_kernel,
@@ -171,7 +171,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.convert.device,
             dim_range=(4, 4),
@@ -191,4 +191,4 @@ def forward(
         # HaarPSI
         l = haarpsi(input, target, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/lpips.py

@@ -18,7 +18,7 @@
 import torchvision.models as models
 import torch.hub as hub
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 
 from typing import Dict, List
 
@@ -207,7 +207,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.shift.device,
             dim_range=(4, 4),
@@ -232,4 +232,4 @@ def forward(
 
         l = torch.stack(residuals, dim=-1).sum(dim=-1)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/mdsi.py

@@ -16,7 +16,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.color import ColorConv
 from piqa.utils.functional import (
     prewitt_kernel,
@@ -178,7 +178,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -202,4 +202,4 @@ def forward(
         # MDSI
         l = mdsi(input, target, kernel=self.kernel, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/psnr.py

@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 
 
 @_jit
@@ -109,7 +109,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=input.device,
             dim_range=(1, -1),
@@ -118,4 +118,4 @@ def forward(
 
         l = psnr(input, target, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/ssim.py

@@ -22,7 +22,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.functional import (
     gaussian_kernel,
     kernel_views,
@@ -251,7 +251,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(3, -1),
@@ -261,7 +261,7 @@ def forward(
 
         l = ssim(input, target, kernel=self.kernel, **self.kwargs)[0]
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)
 
 
 class MS_SSIM(nn.Module):
@@ -331,7 +331,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -347,4 +347,4 @@ def forward(
             **self.kwargs,
         )
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/tv.py

@@ -9,7 +9,7 @@
 import torch
 import torch.nn as nn
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 
 
 @_jit
@@ -94,8 +94,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type([input], device=input.device, dim_range=(3, -1))
+        assert_type([input], device=input.device, dim_range=(3, -1))
 
         l = tv(input, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)

piqa/utils/__init__.py

@@ -13,14 +13,25 @@
     _jit = lambda f: f
 
 
-def _debug(mode: bool = __debug__) -> bool:
-    r"""Returns whether debugging is enabled or not.
+__piqa_debug__ = __debug__
+
+def set_debug(mode: bool = False) -> bool:
+    r"""Sets and returns whether debugging is enabled or not.
+    If `__debug__` is `False`, this function has not effect.
+
+    Example:
+        >>> set_debug(False)
+        False
     """
 
-    return mode
+    global __piqa_debug__
+
+    __piqa_debug__ = __debug__ and mode
+
+    return __piqa_debug__
 
 
-def _assert_type(
+def assert_type(
     tensors: List[torch.Tensor],
     device: torch.device,
     dim_range: Tuple[int, int] = (0, -1),
@@ -33,60 +44,60 @@ def _assert_type(
     Example:
         >>> x = torch.rand(5, 3, 256, 256)
         >>> y = torch.rand(5, 3, 256, 256)
-        >>> _assert_type([x, y], device=x.device, dim_range=(4, 4), n_channels=3)
+        >>> assert_type([x, y], device=x.device, dim_range=(4, 4), n_channels=3)
     """
 
-    if not _debug():
+    if not __piqa_debug__:
         return
 
     ref = tensors[0]
 
     for t in tensors:
         assert t.device == device, (
-            f'Expected tensors to be on {device}, got {t.device}'
+            f'Tensors expected to be on {device}, got {t.device}'
         )
 
         assert t.shape == ref.shape, (
-            'Expected tensors to be of the same shape, got'
+            'Tensors expected to be of the same shape, got'
             f' {ref.shape} and {t.shape}'
         )
 
         if dim_range[0] == dim_range[1]:
             assert t.dim() == dim_range[0], (
-                'Expected number of dimensions to be'
+                'Number of dimensions expected to be'
                 f' {dim_range[0]}, got {t.dim()}'
             )
         elif dim_range[0] < dim_range[1]:
             assert dim_range[0] <= t.dim() <= dim_range[1], (
-                'Expected number of dimensions to be between'
+                'Number of dimensions expected to be between'
                 f' {dim_range[0]} and {dim_range[1]}, got {t.dim()}'
             )
         elif dim_range[0] > 0:
             assert dim_range[0] <= t.dim(), (
-                'Expected number of dimensions to be greater or equal to'
+                'Number of dimensions expected to be greater or equal to'
                 f' {dim_range[0]}, got {t.dim()}'
             )
 
         if n_channels > 0:
             assert t.size(1) == n_channels, (
-                'Expected number of channels to be'
+                'Number of channels expected to be'
                 f' {n_channels}, got {t.size(1)}'
             )
 
         if value_range[0] < value_range[1]:
             assert value_range[0] <= t.min(), (
-                'Expected values to be greater or equal to'
+                'Values expected to be greater or equal to'
                 f' {value_range[0]}, got {t.min()}'
             )
 
             assert t.max() <= value_range[1], (
-                'Expected values to be lower or equal to'
+                'Values expected to be lower or equal to'
                 f' {value_range[1]}, got {t.max()}'
             )
 
 
 @_jit
-def _reduce(x: torch.Tensor, reduction: str = 'mean') -> torch.Tensor:
+def reduce_tensor(x: torch.Tensor, reduction: str = 'mean') -> torch.Tensor:
     r"""Returns the reduction of \(x\).
 
     Args:
@@ -96,7 +107,7 @@ def _reduce(x: torch.Tensor, reduction: str = 'mean') -> torch.Tensor:
 
     Example:
         >>> x = torch.arange(5)
-        >>> _reduce(x, reduction='sum')
+        >>> reduce_tensor(x, reduction='sum')
         tensor(10)
     """
 

piqa/utils/color.py

@@ -83,6 +83,9 @@ def device(self) -> torch.device:
         return self.weight.device
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Defines the computation performed at every call.
+        """
+
         return color_conv(x, self.weight)
 
 

piqa/vsi.py

@@ -23,7 +23,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from piqa.utils import _jit, _assert_type, _reduce
+from piqa.utils import _jit, assert_type, reduce_tensor
 from piqa.utils.color import ColorConv, rgb_to_xyz, xyz_to_lab
 from piqa.utils.functional import (
     scharr_kernel,
@@ -261,7 +261,7 @@ def forward(
         r"""Defines the computation performed at every call.
         """
 
-        _assert_type(
+        assert_type(
             [input, target],
             device=self.kernel.device,
             dim_range=(4, 4),
@@ -292,4 +292,4 @@ def forward(
         # VSI
         l = vsi(input, target, vs_input, vs_target, kernel=self.kernel, **self.kwargs)
 
-        return _reduce(l, self.reduction)
+        return reduce_tensor(l, self.reduction)