Merge branch 'master' into forward_and_compute_to_device

Author: mergify[bot]

torchmetrics/functional/image/psnr.py

@@ -26,6 +26,31 @@ def _psnr_compute(
     base: float = 10.0,
     reduction: str = 'elementwise_mean',
 ) -> Tensor:
+    """
+    Computes peak signal-to-noise ratio.
+
+    Args:
+        sum_squared_error: Sum of square of errors over all observations
+        n_obs: Number of predictions or observations
+        data_range:
+            the range of the data. If None, it is determined from the data (max - min). ``data_range`` must be given
+            when ``dim`` is not None.
+        base: a base of a logarithm to use (default: 10)
+        reduction: a method to reduce metric score over labels.
+
+            - ``'elementwise_mean'``: takes the mean (default)
+            - ``'sum'``: takes the sum
+            - ``'none'``: no reduction will be applied
+
+    Example:
+        >>> preds = torch.tensor([[0.0, 1.0], [2.0, 3.0]])
+        >>> target = torch.tensor([[3.0, 2.0], [1.0, 0.0]])
+        >>> data_range = target.max() - target.min()
+        >>> sum_squared_error, n_obs = _psnr_update(preds, target)
+        >>> _psnr_compute(sum_squared_error, n_obs, data_range)
+        tensor(2.5527)
+    """
+
     psnr_base_e = 2 * torch.log(data_range) - torch.log(sum_squared_error / n_obs)
     psnr_vals = psnr_base_e * (10 / torch.log(tensor(base)))
     return reduce(psnr_vals, reduction=reduction)
@@ -36,6 +61,17 @@ def _psnr_update(
     target: Tensor,
     dim: Optional[Union[int, Tuple[int, ...]]] = None,
 ) -> Tuple[Tensor, Tensor]:
+    """
+    Updates and returns variables required to compute peak signal-to-noise ratio.
+
+    Args:
+        preds: Predicted tensor
+        target: Ground truth tensor
+        dim:
+            Dimensions to reduce PSNR scores over provided as either an integer or a list of integers. Default is
+            None meaning scores will be reduced across all dimensions.
+    """
+
     if dim is None:
         sum_squared_error = torch.sum(torch.pow(preds - target, 2))
         n_obs = tensor(target.numel(), device=target.device)
@@ -66,7 +102,7 @@ def psnr(
     dim: Optional[Union[int, Tuple[int, ...]]] = None,
 ) -> Tensor:
     """
-    Computes the peak signal-to-noise ratio
+    Computes the peak signal-to-noise ratio.
 
     Args:
         preds: estimated signal

torchmetrics/functional/image/ssim.py

@@ -22,6 +22,19 @@
 
 
 def _gaussian(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.device) -> Tensor:
+    """
+    Computes 1D gaussian kernel
+
+    Args:
+        kernel_size: size of the gaussian kernel
+        sigma: Standard deviation of the gaussian kernel
+        dtype: data type of the output tensor
+        device: device of the output tensor
+
+    Example:
+        >>> _gaussian(3, 1, torch.float, 'cpu')
+        tensor([[0.2741, 0.4519, 0.2741]])
+    """
     dist = torch.arange(start=(1 - kernel_size) / 2, end=(1 + kernel_size) / 2, step=1, dtype=dtype, device=device)
     gauss = torch.exp(-torch.pow(dist / sigma, 2) / 2)
     return (gauss / gauss.sum()).unsqueeze(dim=0)  # (1, kernel_size)
@@ -30,6 +43,25 @@ def _gaussian(kernel_size: int, sigma: float, dtype: torch.dtype, device: torch.
 def _gaussian_kernel(
     channel: int, kernel_size: Sequence[int], sigma: Sequence[float], dtype: torch.dtype, device: torch.device
 ) -> Tensor:
+    """
+    Computes 2D gaussian kernel
+
+    Args:
+        channel: number of channels in the image
+        kernel_size: size of the gaussian kernel as a tuple (h, w)
+        sigma: Standard deviation of the gaussian kernel
+        dtype: data type of the output tensor
+        device: device of the output tensor
+
+    Example:
+        >>> _gaussian_kernel(1, (5,5), (1,1), torch.float, "cpu")
+        tensor([[[[0.0030, 0.0133, 0.0219, 0.0133, 0.0030],
+                  [0.0133, 0.0596, 0.0983, 0.0596, 0.0133],
+                  [0.0219, 0.0983, 0.1621, 0.0983, 0.0219],
+                  [0.0133, 0.0596, 0.0983, 0.0596, 0.0133],
+                  [0.0030, 0.0133, 0.0219, 0.0133, 0.0030]]]])
+    """
+
     gaussian_kernel_x = _gaussian(kernel_size[0], sigma[0], dtype, device)
     gaussian_kernel_y = _gaussian(kernel_size[1], sigma[1], dtype, device)
     kernel = torch.matmul(gaussian_kernel_x.t(), gaussian_kernel_y)  # (kernel_size, 1) * (1, kernel_size)
@@ -38,6 +70,15 @@ def _gaussian_kernel(
 
 
 def _ssim_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
+    """
+    Updates and returns variables required to compute Structural Similarity Index Measure.
+    Checks for same shape and type of the input tensors.
+
+    Args:
+        preds: Predicted tensor
+        target: Ground truth tensor
+    """
+
     if preds.dtype != target.dtype:
         raise TypeError(
             "Expected `preds` and `target` to have the same data type."
@@ -62,6 +103,31 @@ def _ssim_compute(
     k1: float = 0.01,
     k2: float = 0.03,
 ) -> Tensor:
+    """
+    Computes Structual Similarity Index Measure
+
+    Args:
+        preds: estimated image
+        target: ground truth image
+        kernel_size: size of the gaussian kernel (default: (11, 11))
+        sigma: Standard deviation of the gaussian kernel (default: (1.5, 1.5))
+        reduction: a method to reduce metric score over labels.
+
+            - ``'elementwise_mean'``: takes the mean (default)
+            - ``'sum'``: takes the sum
+            - ``'none'``: no reduction will be applied
+
+        data_range: Range of the image. If ``None``, it is determined from the image (max - min)
+        k1: Parameter of SSIM. Default: 0.01
+        k2: Parameter of SSIM. Default: 0.03
+
+    Example:
+        >>> preds = torch.rand([16, 1, 16, 16])
+        >>> target = preds * 0.75
+        >>> preds, target = _ssim_update(preds, target)
+        >>> _ssim_compute(preds, target)
+        tensor(0.9219)
+    """
     if len(kernel_size) != 2 or len(sigma) != 2:
         raise ValueError(
             "Expected `kernel_size` and `sigma` to have the length of two."

torchmetrics/functional/retrieval/ndcg.py

@@ -20,6 +20,7 @@
 
 
 def _dcg(target: Tensor) -> Tensor:
+    """ Computes Discounted Cumulative Gain for input tensor """
     denom = torch.log2(torch.arange(target.shape[-1], device=target.device) + 2.0)
     return (target / denom).sum(dim=-1)
 

torchmetrics/functional/text/bleu.py

@@ -54,6 +54,19 @@ def _bleu_score_update(
     ref_len: Tensor,
     n_gram: int = 4
 ) -> Tuple[Tensor, Tensor]:
+    """
+    Updates and returns variables required to compute the BLEU score.
+
+    Args:
+        reference_corpus: An iterable of iterables of reference corpus
+        translate_corpus: An iterable of machine translated corpus
+        numerator: Numerator of precision score (true positives)
+        denominator: Denominator of precision score (true positives + false positives)
+        trans_len: count of words in a candidate translation
+        ref_len: count of words in a reference translation
+        n_gram: gram value ranged 1 to 4
+    """
+
     for (translation, references) in zip(translate_corpus, reference_corpus):
         trans_len += len(translation)
         ref_len_list = [len(ref) for ref in references]
@@ -84,6 +97,17 @@ def _bleu_score_compute(
     n_gram: int = 4,
     smooth: bool = False
 ) -> Tensor:
+    """
+    Computes the BLEU score.
+
+    Args:
+        trans_len: count of words in a candidate translation
+        ref_len: count of words in a reference translation
+        numerator: Numerator of precision score (true positives)
+        denominator: Denominator of precision score (true positives + false positives)
+        n_gram: gram value ranged 1 to 4
+        smooth: Whether or not to apply smoothing
+    """
     if min(numerator) == 0.0:
         return tensor(0.0)