Add PIT for audio metrics

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added Symmetric Mean Absolute Percentage error (SMAPE) ([#375](https://github.com/PyTorchLightning/metrics/issues/375))
 
 
+- Added Permutation Invariant Training metric (PIT) ([#294](https://github.com/PyTorchLightning/metrics/issues/294))
+
+
 ### Changed
 
 - Moved `psnr` and `ssim` from `torchmetrics.functional.regression.*` to `torchmetrics.functional.image.*` ([#382](https://github.com/PyTorchLightning/metrics/pull/382))

docs/source/references/functional.rst

@@ -9,6 +9,13 @@ Functional metrics
 Audio Metrics
 *************
 
+pit [func]
+~~~~~~~~~~
+
+.. autofunction:: torchmetrics.functional.pit
+    :noindex:
+
+
 si_sdr [func]
 ~~~~~~~~~~~~~
 

docs/source/references/modules.rst

@@ -40,6 +40,12 @@ the metric will be computed over the ``time`` dimension.
     >>> snr_val
     tensor(16.1805)
 
+PIT
+~~~~~~
+
+.. autoclass:: torchmetrics.PIT
+    :noindex:
+
 SI_SDR
 ~~~~~~
 

tests/audio/test_pit.py

@@ -0,0 +1,210 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from collections import namedtuple
+from functools import partial
+from typing import Callable, Tuple
+
+import numpy as np
+import pytest
+import torch
+from scipy.optimize import linear_sum_assignment
+from torch import Tensor
+
+from tests.helpers import seed_all
+from tests.helpers.testers import BATCH_SIZE, NUM_BATCHES, MetricTester
+from torchmetrics.audio import PIT
+from torchmetrics.functional import pit, si_sdr, snr
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_6
+
+seed_all(42)
+
+TIME = 10
+
+Input = namedtuple('Input', ["preds", "target"])
+
+# three speaker examples to test _find_best_perm_by_linear_sum_assignment
+inputs1 = Input(
+    preds=torch.rand(NUM_BATCHES, BATCH_SIZE, 3, TIME),
+    target=torch.rand(NUM_BATCHES, BATCH_SIZE, 3, TIME),
+)
+# two speaker examples to test _find_best_perm_by_exhuastive_method
+inputs2 = Input(
+    preds=torch.rand(NUM_BATCHES, BATCH_SIZE, 2, TIME),
+    target=torch.rand(NUM_BATCHES, BATCH_SIZE, 2, TIME),
+)
+
+
+def a_naive_implementation_of_pit_based_on_scipy(
+    preds: Tensor,
+    target: Tensor,
+    metric_func: Callable,
+    eval_func: str,
+) -> Tuple[Tensor, Tensor]:
+    """a naive implementation of pit based on scipy
+
+    Args:
+        preds: predictions, shape[batch, spk, time]
+        target: targets, shape[batch, spk, time]
+        metric_func: which metric
+        eval_func: min or max
+
+    Returns:
+        best_metric:
+            shape [batch]
+        best_perm:
+            shape [batch, spk]
+    """
+    batch_size, spk_num = target.shape[0:2]
+    metric_mtx = torch.empty((batch_size, spk_num, spk_num), device=target.device)
+    for t in range(spk_num):
+        for e in range(spk_num):
+            metric_mtx[:, t, e] = metric_func(preds[:, e, ...], target[:, t, ...])
+
+    # pit_r = PIT(metric_func, eval_func)(preds, target)
+    metric_mtx = metric_mtx.detach().cpu().numpy()
+    best_metrics = []
+    best_perms = []
+    for b in range(batch_size):
+        row_idx, col_idx = linear_sum_assignment(metric_mtx[b, ...], eval_func == 'max')
+        best_metrics.append(metric_mtx[b, row_idx, col_idx].mean())
+        best_perms.append(col_idx)
+    return torch.from_numpy(np.stack(best_metrics)), torch.from_numpy(np.stack(best_perms))
+
+
+def average_metric(preds: Tensor, target: Tensor, metric_func: Callable) -> Tensor:
+    """average the metric values
+
+    Args:
+        preds: predictions, shape[batch, spk, time]
+        target: targets, shape[batch, spk, time]
+        metric_func: a function which return best_metric and best_perm
+
+    Returns:
+        the average of best_metric
+    """
+    return metric_func(preds, target)[0].mean()
+
+
+snr_pit_scipy = partial(a_naive_implementation_of_pit_based_on_scipy, metric_func=snr, eval_func='max')
+si_sdr_pit_scipy = partial(a_naive_implementation_of_pit_based_on_scipy, metric_func=si_sdr, eval_func='max')
+
+
+@pytest.mark.parametrize(
+    "preds, target, sk_metric, metric_func, eval_func",
+    [
+        (inputs1.preds, inputs1.target, snr_pit_scipy, snr, 'max'),
+        (inputs1.preds, inputs1.target, si_sdr_pit_scipy, si_sdr, 'max'),
+        (inputs2.preds, inputs2.target, snr_pit_scipy, snr, 'max'),
+        (inputs2.preds, inputs2.target, si_sdr_pit_scipy, si_sdr, 'max'),
+    ],
+)
+class TestPIT(MetricTester):
+    atol = 1e-2
+
+    @pytest.mark.parametrize("ddp", [True, False])
+    @pytest.mark.parametrize("dist_sync_on_step", [True, False])
+    def test_pit(self, preds, target, sk_metric, metric_func, eval_func, ddp, dist_sync_on_step):
+        self.run_class_metric_test(
+            ddp,
+            preds,
+            target,
+            PIT,
+            sk_metric=partial(average_metric, metric_func=sk_metric),
+            dist_sync_on_step=dist_sync_on_step,
+            metric_args=dict(metric_func=metric_func, eval_func=eval_func),
+        )
+
+    def test_pit_functional(self, preds, target, sk_metric, metric_func, eval_func):
+        device = 'cuda' if (torch.cuda.is_available() and torch.cuda.device_count() > 0) else 'cpu'
+
+        # move to device
+        preds = preds.to(device)
+        target = target.to(device)
+
+        for i in range(NUM_BATCHES):
+            best_metric, best_perm = pit(preds[i], target[i], metric_func, eval_func)
+            best_metric_sk, best_perm_sk = sk_metric(preds[i].cpu(), target[i].cpu())
+
+            # assert its the same
+            assert np.allclose(
+                best_metric.detach().cpu().numpy(), best_metric_sk.detach().cpu().numpy(), atol=self.atol
+            )
+            assert (best_perm.detach().cpu().numpy() == best_perm_sk.detach().cpu().numpy()).all()
+
+    def test_pit_differentiability(self, preds, target, sk_metric, metric_func, eval_func):
+
+        def pit_diff(preds, target, metric_func, eval_func):
+            return pit(preds, target, metric_func, eval_func)[0]
+
+        self.run_differentiability_test(
+            preds=preds,
+            target=target,
+            metric_module=PIT,
+            metric_functional=pit_diff,
+            metric_args={
+                'metric_func': metric_func,
+                'eval_func': eval_func
+            }
+        )
+
+    @pytest.mark.skipif(
+        not _TORCH_GREATER_EQUAL_1_6, reason='half support of core operations on not support before pytorch v1.6'
+    )
+    def test_pit_half_cpu(self, preds, target, sk_metric, metric_func, eval_func):
+        pytest.xfail("PIT metric does not support cpu + half precision")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='test requires cuda')
+    def test_pit_half_gpu(self, preds, target, sk_metric, metric_func, eval_func):
+        self.run_precision_test_gpu(
+            preds=preds,
+            target=target,
+            metric_module=PIT,
+            metric_functional=partial(pit, metric_func=metric_func, eval_func=eval_func),
+            metric_args={
+                'metric_func': metric_func,
+                'eval_func': eval_func
+            }
+        )
+
+
+def test_error_on_different_shape() -> None:
+    metric = PIT(snr, 'max')
+    with pytest.raises(RuntimeError, match='Predictions and targets are expected to have the same shape'):
+        metric(torch.randn(3, 3, 10), torch.randn(3, 2, 10))
+
+
+def test_error_on_wrong_eval_func() -> None:
+    metric = PIT(snr, 'xxx')
+    with pytest.raises(RuntimeError, match='eval_func can only be "max" or "min"'):
+        metric(torch.randn(3, 3, 10), torch.randn(3, 3, 10))
+
+
+def test_error_on_wrong_shape() -> None:
+    metric = PIT(snr, 'max')
+    with pytest.raises(RuntimeError, match='Inputs must be of shape *'):
+        metric(torch.randn(3), torch.randn(3))
+
+
+def test_consistency_of_two_implementations() -> None:
+    from torchmetrics.functional.audio.pit import (
+        _find_best_perm_by_exhuastive_method,
+        _find_best_perm_by_linear_sum_assignment,
+    )
+    shapes_test = [(5, 2, 2), (4, 3, 3), (4, 4, 4), (3, 5, 5)]
+    for shp in shapes_test:
+        metric_mtx = torch.randn(size=shp)
+        bm1, bp1 = _find_best_perm_by_linear_sum_assignment(metric_mtx, torch.max)
+        bm2, bp2 = _find_best_perm_by_exhuastive_method(metric_mtx, torch.max)
+        assert torch.allclose(bm1, bm2)
+        assert (bp1 == bp2).all()

torchmetrics/__init__.py

@@ -11,7 +11,7 @@
 _PACKAGE_ROOT = os.path.dirname(__file__)
 _PROJECT_ROOT = os.path.dirname(_PACKAGE_ROOT)
 
-from torchmetrics.audio import SI_SDR, SI_SNR, SNR  # noqa: F401 E402
+from torchmetrics.audio import PIT, SI_SDR, SI_SNR, SNR  # noqa: F401 E402
 from torchmetrics.average import AverageMeter  # noqa: F401 E402
 from torchmetrics.classification import (  # noqa: F401 E402
     AUC,

torchmetrics/audio/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from torchmetrics.audio.pit import PIT  # noqa: F401
 from torchmetrics.audio.si_sdr import SI_SDR  # noqa: F401
 from torchmetrics.audio.si_snr import SI_SNR  # noqa: F401
 from torchmetrics.audio.snr import SNR  # noqa: F401

torchmetrics/audio/pit.py

@@ -0,0 +1,114 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Callable, Dict, Optional
+
+from torch import Tensor, tensor
+
+from torchmetrics.functional.audio.pit import pit
+from torchmetrics.metric import Metric
+
+
+class PIT(Metric):
+    """ Permutation invariant training metric
+
+    Forward accepts
+
+    - ``preds``: ``shape [..., time]``
+    - ``target``: ``shape [..., time]``
+
+    Args:
+        metric_func:
+            a metric function accept a batch of target and estimate, i.e. metric_func(target[:, i, ...],
+            estimate[:, j, ...]), and returns a batch of metric tensors [batch]
+        eval_func:
+            the function to find the best permutation, can be 'min' or 'max', i.e. the smaller the better
+            or the larger the better.
+        compute_on_step:
+            Forward only calls ``update()`` and returns None if this is set to False. default: True
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: None (which selects the entire world)
+        dist_sync_fn:
+            Callback that performs the allgather operation on the metric state. When `None`, DDP
+            will be used to perform the allgather.
+        kwargs:
+            additional args for metric_func
+
+    Returns:
+        average PIT metric
+
+    Example:
+        >>> import torch
+        >>> from torchmetrics import PIT
+        >>> from torchmetrics.functional.audio import si_snr
+        >>> preds = torch.randn(3, 2, 5) # [batch, spk, time]
+        >>> target = torch.randn(3, 2, 5) # [batch, spk, time]
+        >>> pit = PIT(si_snr, 'max')
+        >>> avg_pit_metric = pit(preds, target)
+
+    Reference:
+        [1]	D. Yu, M. Kolbaek, Z.-H. Tan, J. Jensen, Permutation invariant training of deep models for
+        speaker-independent multi-talker speech separation, in: 2017 IEEE Int. Conf. Acoust. Speech
+        Signal Process. ICASSP, IEEE, New Orleans, LA, 2017: pp. 241–245. https://doi.org/10.1109/ICASSP.2017.7952154.
+    """
+    sum_pit_metric: Tensor
+    total: Tensor
+
+    def __init__(
+        self,
+        metric_func: Callable,
+        eval_func: str = 'max',
+        compute_on_step: bool = True,
+        dist_sync_on_step: bool = False,
+        process_group: Optional[Any] = None,
+        dist_sync_fn: Optional[Callable[[Tensor], Tensor]] = None,
+        **kwargs: Dict[str, Any],
+    ) -> None:
+        super().__init__(
+            compute_on_step=compute_on_step,
+            dist_sync_on_step=dist_sync_on_step,
+            process_group=process_group,
+            dist_sync_fn=dist_sync_fn,
+        )
+        self.metric_func = metric_func
+        self.eval_func = eval_func
+        self.kwargs = kwargs
+
+        self.add_state("sum_pit_metric", default=tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
+
+    def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
+        """
+        Update state with predictions and targets.
+
+        Args:
+            preds: Predictions from model
+            target: Ground truth values
+        """
+        pit_metric = pit(preds, target, self.metric_func, self.eval_func, **self.kwargs)[0]
+
+        self.sum_pit_metric += pit_metric.sum()
+        self.total += pit_metric.numel()
+
+    def compute(self) -> Tensor:
+        """
+        Computes average PIT metric.
+        """
+        return self.sum_pit_metric / self.total
+
+    @property
+    def is_differentiable(self) -> bool:
+        return True

torchmetrics/functional/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from torchmetrics.functional.audio.pit import permutate, pit  # noqa: F401
 from torchmetrics.functional.audio.si_sdr import si_sdr  # noqa: F401
 from torchmetrics.functional.audio.si_snr import si_snr  # noqa: F401
 from torchmetrics.functional.audio.snr import snr  # noqa: F401

torchmetrics/functional/audio/__init__.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from torchmetrics.functional.audio.pit import permutate, pit  # noqa: F401
 from torchmetrics.functional.audio.si_sdr import si_sdr  # noqa: F401
 from torchmetrics.functional.audio.si_snr import si_snr  # noqa: F401
 from torchmetrics.functional.audio.snr import snr  # noqa: F401