Total Variation Loss

torchmetrics/functional/image/__init__.py

@@ -20,4 +20,5 @@
     multiscale_structural_similarity_index_measure,
     structural_similarity_index_measure,
 )
+from torchmetrics.functional.image.total_variation import total_variation  # noqa: F401
 from torchmetrics.functional.image.uqi import universal_image_quality_index  # noqa: F401

torchmetrics/functional/image/total_variation.py

@@ -0,0 +1,24 @@
+import torch
+
+
+def total_variation(img: torch.Tensor) -> torch.Tensor:
+    """Computes total variation loss.
+
+    Adapted from https://github.com/jxgu1016/Total_Variation_Loss.pytorch
+    Args:
+        img (torch.Tensor): A NCHW image batch.
+
+    Returns:
+        A loss scalar value.
+    """
+
+    def tensor_size(t):
+        return t.size()[1] * t.size()[2] * t.size()[3]
+
+    _height = img.size()[2]
+    _width = img.size()[3]
+    _count_height = tensor_size(img[:, :, 1:, :])
+    _count_width = tensor_size(img[:, :, :, 1:])
+    _height_tv = torch.pow((img[:, :, 1:, :] - img[:, :, : _height - 1, :]), 2).sum()
+    _width_tv = torch.pow((img[:, :, :, 1:] - img[:, :, :, : _width - 1]), 2).sum()
+    return (2 * (_height_tv / _count_height + _width_tv / _count_width)) / img.size()[0]

torchmetrics/image/__init__.py

@@ -29,3 +29,5 @@
 
 if _LPIPS_AVAILABLE:
     from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity  # noqa: F401
+
+from torchmetrics.image.total_variation import TotalVariation  # noqa: F401

torchmetrics/image/total_variation.py

@@ -0,0 +1,50 @@
+import torch
+
+from torchmetrics.metric import Metric
+
+
+class TotalVariation(Metric):
+    """Computes Total Variation loss.
+
+    Adapted from: https://github.com/jxgu1016/Total_Variation_Loss.pytorch
+    Args:
+        dist_sync_on_step: Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        compute_on_step: Forward only calls ``update()`` and returns None if this is set to
+            False.
+    """
+
+    is_differentiable = True
+    higher_is_better = False
+    current: torch.Tensor
+    total: torch.Tensor
+
+    def __init__(self, dist_sync_on_step: bool = False, compute_on_step: bool = True):
+        super().__init__(dist_sync_on_step=dist_sync_on_step, compute_on_step=compute_on_step)
+        self.add_state("current", default=torch.tensor(0, dtype=torch.float), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0, dtype=torch.int), dist_reduce_fx="sum")
+
+    def update(self, img: torch.Tensor) -> None:
+        """Update method for TV Loss.
+
+        Args:
+            img (torch.Tensor): A NCHW image batch.
+
+        Returns:
+            A loss scalar value.
+        """
+        _height = img.size()[2]
+        _width = img.size()[3]
+        _count_height = self.tensor_size(img[:, :, 1:, :])
+        _count_width = self.tensor_size(img[:, :, :, 1:])
+        _height_tv = torch.pow((img[:, :, 1:, :] - img[:, :, : _height - 1, :]), 2).sum()
+        _width_tv = torch.pow((img[:, :, :, 1:] - img[:, :, :, : _width - 1]), 2).sum()
+        self.current += 2 * (_height_tv / _count_height + _width_tv / _count_width)
+        self.total += img.numel()
+
+    def compute(self):
+        return self.current.float() / self.total
+
+    @staticmethod
+    def tensor_size(t):
+        return t.size()[1] * t.size()[2] * t.size()[3]