Repository: Lightning-AI/torchmetrics

Issue:
The sdr metric in TM sometime gives NaN for some input
## 🐛 Bug

This issue is related with `fast-bss-eval`'s torch version, see https://github.com/fakufaku/fast_bss_eval/issues/5

### To Reproduce

```
import numpy as np
import torch

x = np.load('debug.npz')
preds = torch.tensor(x['preds'])
target = torch.tensor(x['target'])
print(preds.shape, target.shape)

from torchmetrics.functional.audio import signal_distortion_ratio
sdr = signal_distortion_ratio(preds, target)
print(sdr)

from mir_eval.separation import bss_eval_sources
sdr, _, _, _ = bss_eval_sources(target.numpy(), preds.numpy(), False)
print(sdr)
```

outputs:
```
torch.Size([2, 64000]) torch.Size([2, 64000])
tensor([-2.6815,     nan])
[-2.68156071 44.58523729]
```

unzip [data.zip](https://github.com/PyTorchLightning/metrics/files/8309638/data.zip) to get the `debug.npz`


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

#### Code sample

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

the results given by signal_distortion_ratio is close to the one given by `mir_eval`
<!-- A clear and concise description of what you expected to happen. -->

### Environment

- OS (e.g., Linux):
- Python & PyTorch Version (e.g., 1.0):
- How you installed PyTorch (`conda`, `pip`, build command if you used source):
- Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi @quancs,
From the issue that you linked it seems that the solution from the author is basically to do the evaluation in `double` instead of `float`. I can confirm that doing this fixes the example you send. Do you think we should cast the users input here:
https://github.com/PyTorchLightning/metrics/blob/865a08fcf102c2eb1b776b13643bc87aadf7f4f7/torchmetrics/functional/audio/sdr.py#L140-L141
to `double` instead of `float`. Alternatively, we can insert note in docstring that in some cases it is better to evaluate using double precision.
> Hi @quancs, From the issue that you linked it seems that the solution from the author is basically to do the evaluation in `double` instead of `float`. I can confirm that doing this fixes the example you send. Do you think we should cast the users input here:
> 
> https://github.com/PyTorchLightning/metrics/blob/865a08fcf102c2eb1b776b13643bc87aadf7f4f7/torchmetrics/functional/audio/sdr.py#L140-L141
> 
> to `double` instead of `float`. Alternatively, we can insert note in docstring that in some cases it is better to evaluate using double precision.

This issue happens on the torch version on cpu. On GPU it's OK.
And I don't see any vialation in my past experiment results tested on GPU.
I have some ideas to fix this:
1) convert to double anyway, no matter on CPU or GPU, but it may make the metric slow
2) convert to double on CPU, but it may make the metric slow, and we don't know whether on GPU is really OK
3) convert to double when we detect the result is not a valid number (NaN or Inf), and run again.

I prefer 3). what do you think? or do you have other ideas?
3) is fine by me :]

Resolving patch:
diff --git a/torchmetrics/functional/audio/sdr.py b/torchmetrics/functional/audio/sdr.py
--- a/torchmetrics/functional/audio/sdr.py
+++ b/torchmetrics/functional/audio/sdr.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from typing import Optional
 
 import torch
@@ -192,6 +193,15 @@ def signal_distortion_ratio(
     # transform to decibels
     ratio = coh / (1 - coh)
     val = 10.0 * torch.log10(ratio)
+
+    # recompute sdr in float64 if val is NaN or Inf
+    if (torch.isnan(val).any() or torch.isinf(val).any()) and preds.dtype != torch.float64:
+        warnings.warn(
+            "Detected `nan` or `inf` value in computed metric, retrying computation in double precision",
+            UserWarning,
+        )
+        val = signal_distortion_ratio(preds.double(), target.double(), use_cg_iter, filter_length, zero_mean, load_diag)
+
     return val
 
 


Test patch:
diff --git a/tests/audio/test_sdr.py b/tests/audio/test_sdr.py
--- a/tests/audio/test_sdr.py
+++ b/tests/audio/test_sdr.py
@@ -15,13 +15,14 @@
 from functools import partial
 from typing import Callable
 
+import numpy as np
 import pytest
 import torch
 from mir_eval.separation import bss_eval_sources
 from scipy.io import wavfile
 from torch import Tensor
 
-from tests.audio import _SAMPLE_AUDIO_SPEECH, _SAMPLE_AUDIO_SPEECH_BAB_DB
+from tests.audio import _SAMPLE_AUDIO_SPEECH, _SAMPLE_AUDIO_SPEECH_BAB_DB, _SAMPLE_NUMPY_ISSUE_895
 from tests.helpers import seed_all
 from tests.helpers.testers import MetricTester
 from torchmetrics.audio import SignalDistortionRatio
@@ -146,3 +147,24 @@ def test_on_real_audio():
         rtol=0.0001,
         atol=1e-4,
     )
+
+
+@pytest.mark.skipif(
+    not _TORCH_GREATER_EQUAL_1_8, reason="when pytorch < 1.8, sdr is using numpy which doesn't have this problem"
+)
+def test_too_low_precision():
+    """Corner case where the precision of the input is important."""
+    data = np.load(_SAMPLE_NUMPY_ISSUE_895)
+    preds = torch.tensor(data["preds"])
+    target = torch.tensor(data["target"])
+    with pytest.warns(
+        UserWarning, match="Detected `nan` or `inf` value in computed metric, retrying computation in double precision"
+    ):
+        sdr_tm = signal_distortion_ratio(preds, target)
+    sdr_bss, _, _, _ = bss_eval_sources(target.numpy(), preds.numpy(), False)
+    assert torch.allclose(
+        sdr_tm.mean(),
+        torch.tensor(sdr_bss).mean(),
+        rtol=0.0001,
+        atol=1e-4,
+    )