Repository: pytorch/audio

Issue:
torchaudio.transforms.Spectrogram with power=1 can easily NaN
## 🐛 Bug

If you use STFT through Spectrogram, with power=1, you can easily get a NaN. If your spectrogram square magnitude underflows to zero, you end up taking backprop of sqrt(0), which leads to a NaN. 

## To Reproduce

Steps to reproduce the behavior:

1. Use torchaudio.transforms.Spectrogram, set power=1, do a backward pass.

## Expected behavior

There should be a configurable epsilon to clamp the input to sqrt, otherwise users have to use torch.stft directly or use other workarounds (such as power=2).


Discussion:
Hi @gibiansky 

Thanks for the report. I think I got the gist, but could you illustrate the fix?

Is it here?

https://github.com/pytorch/audio/blob/0076ab073d1ee6160efbc239e075196b35ed850b/torchaudio/functional.py#L95
Not quite. 

Here is the issue. Using torchaudio.transforms.Spectrogram with power=1 is equivalent to:

```
stft = torch.sftf(waveform, n_fft, hop_length, win_length, ...)
norm_sq = stft.pow(2.).sum(-1)
result = norm_sq.sqrt()
```

Let's suppose your waveform is zero, and so your norm_sq is zero. Then you take the square root of that, and you get zero. _However_, during the backward pass, the sqrt of zero is mathematically undefined at one, and torch.sqrt gradient is NaN.

I want to be able to do this:

```
stft = torch.sftf(waveform, n_fft, hop_length, win_length, ...)
norm_sq = stft.pow(2.).sum(-1)
result = (1e-9 + norm_sq).sqrt()
```

This ensures that we cannot get NaN on backward pass.

Here is a related issue also caused by sqrt of zero backward pass: 

https://github.com/pytorch/pytorch/issues/43211
https://github.com/pytorch/pytorch/issues/37323

I think perhaps this has been fixed in `norm()`, but computing the norm by hand still causes the issue, because it wasn't fixed for `sqrt`:

https://github.com/pytorch/pytorch/blob/201ba139115adeecae4f094a9c9790200e53ff99/tools/autograd/templates/Functions.cpp#L115
In fact, the issue is caused by the line [below the one you linked](https://github.com/pytorch/audio/blob/0076ab073d1ee6160efbc239e075196b35ed850b/torchaudio/functional.py#L97), @mthrok. Due to the [`torch.norm` bug linked in the code](https://github.com/pytorch/pytorch/issues/34279), you implemented `complex_norm`, but complex_norm uses sqrt directly, and so it has poor backward pass behaviour (NaN) when the input is zero (unlike torch.norm, which special-cases zero inputs). 

With regards to fixing it, I can see that it might be a bit tough; the core issue is torch.sqrt and torch.norm. torch.norm behaves well on backward pass but is slow on CPU compared to implementing norm yourself. But the implemented complex_norm returns NaN instead of a zero subgradient at input zero, which causes problems.

My original suggestion, adding a tiny epsilon, isn't great either, because it's not really doing the right mathematical operation anymore and can fail for float16 (as in one of the things I linked).

I would recommend (a) switching complex_norm to check the device, and use torch.norm if we're on GPU and (b) changing complex_norm to torch.norm ASAP.
@gibiansky Thanks for the detailed explanation.

@anjali411 If we migrate to ComplexTensor type, can the backward pass of `abs()` handle zero <-> NaN issue mentioned above??
> @anjali411 If we migrate to ComplexTensor type, can the backward pass of `abs()` handle zero <-> NaN issue mentioned above??

yeah here's a sample code to show that:
```
>>> x=torch.zeros(2, 3, dtype=torch.cfloat, requires_grad=True)
>>> y=x.abs()
>>> y.sum().backward()
>>> x.grad
tensor([[0.+0.j, 0.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j, 0.+0.j]])

>>> x1=torch.zeros(2, 3, dtype=torch.float, requires_grad=True)
>>> y1=x1.sqrt()
>>> y1.sum().backward()
>>> x1.grad
tensor([[inf, inf, inf],
        [inf, inf, inf]])
```
@anjali411 Thanks!

@gibiansky Do you think this approach is sufficient for the issue you are facing?
I am unfortunately somewhat new to PyTorch (I have been using it for all of three days at this point) so I may be missing subtleties, but at least on the surface level it seems like this would solve my issue neatly.

Resolving patch:
diff --git a/torchaudio/functional/functional.py b/torchaudio/functional/functional.py
--- a/torchaudio/functional/functional.py
+++ b/torchaudio/functional/functional.py
@@ -70,30 +70,29 @@ def spectrogram(
     waveform = waveform.reshape(-1, shape[-1])
 
     # default values are consistent with librosa.core.spectrum._spectrogram
-    spec_f = torch.view_as_real(
-        torch.stft(
-            input=waveform,
-            n_fft=n_fft,
-            hop_length=hop_length,
-            win_length=win_length,
-            window=window,
-            center=True,
-            pad_mode="reflect",
-            normalized=False,
-            onesided=True,
-            return_complex=True,
-        )
+    spec_f = torch.stft(
+        input=waveform,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=True,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
     )
 
     # unpack batch
-    spec_f = spec_f.reshape(shape[:-1] + spec_f.shape[-3:])
+    spec_f = spec_f.reshape(shape[:-1] + spec_f.shape[-2:])
 
     if normalized:
         spec_f /= window.pow(2.).sum().sqrt()
     if power is not None:
-        spec_f = complex_norm(spec_f, power=power)
-
-    return spec_f
+        if power == 1.0:
+            return spec_f.abs()
+        return spec_f.abs().pow(power)
+    return torch.view_as_real(spec_f)
 
 
 def griffinlim(


Test patch:
diff --git a/test/torchaudio_unittest/functional/functional_cpu_test.py b/test/torchaudio_unittest/functional/functional_cpu_test.py
--- a/test/torchaudio_unittest/functional/functional_cpu_test.py
+++ b/test/torchaudio_unittest/functional/functional_cpu_test.py
@@ -1,5 +1,4 @@
 import math
-import unittest
 
 import torch
 import torchaudio
@@ -8,7 +7,7 @@
 import pytest
 
 from torchaudio_unittest import common_utils
-from .functional_impl import Lfilter
+from .functional_impl import Lfilter, Spectrogram
 
 
 class TestLFilterFloat32(Lfilter, common_utils.PytorchTestCase):
@@ -21,6 +20,16 @@ class TestLFilterFloat64(Lfilter, common_utils.PytorchTestCase):
     device = torch.device('cpu')
 
 
+class TestSpectrogramFloat32(Spectrogram, common_utils.PytorchTestCase):
+    dtype = torch.float32
+    device = torch.device('cpu')
+
+
+class TestSpectrogramFloat64(Spectrogram, common_utils.PytorchTestCase):
+    dtype = torch.float64
+    device = torch.device('cpu')
+
+
 class TestCreateFBMatrix(common_utils.TorchaudioTestCase):
     def test_no_warning_high_n_freq(self):
         with pytest.warns(None) as w:
diff --git a/test/torchaudio_unittest/functional/functional_cuda_test.py b/test/torchaudio_unittest/functional/functional_cuda_test.py
--- a/test/torchaudio_unittest/functional/functional_cuda_test.py
+++ b/test/torchaudio_unittest/functional/functional_cuda_test.py
@@ -1,7 +1,7 @@
 import torch
 
 from torchaudio_unittest import common_utils
-from .functional_impl import Lfilter
+from .functional_impl import Lfilter, Spectrogram
 
 
 @common_utils.skipIfNoCuda
@@ -14,3 +14,15 @@ class TestLFilterFloat32(Lfilter, common_utils.PytorchTestCase):
 class TestLFilterFloat64(Lfilter, common_utils.PytorchTestCase):
     dtype = torch.float64
     device = torch.device('cuda')
+
+
+@common_utils.skipIfNoCuda
+class TestSpectrogramFloat32(Spectrogram, common_utils.PytorchTestCase):
+    dtype = torch.float32
+    device = torch.device('cuda')
+
+
+@common_utils.skipIfNoCuda
+class TestSpectrogramFloat64(Spectrogram, common_utils.PytorchTestCase):
+    dtype = torch.float64
+    device = torch.device('cuda')
diff --git a/test/torchaudio_unittest/functional/functional_impl.py b/test/torchaudio_unittest/functional/functional_impl.py
--- a/test/torchaudio_unittest/functional/functional_impl.py
+++ b/test/torchaudio_unittest/functional/functional_impl.py
@@ -1,6 +1,7 @@
 """Test defintion common to CPU and CUDA"""
 import torch
 import torchaudio.functional as F
+from parameterized import parameterized
 
 from torchaudio_unittest import common_utils
 
@@ -29,3 +30,25 @@ def test_clamp(self):
         assert output_signal.max() <= 1
         output_signal = F.lfilter(input_signal, a_coeffs, b_coeffs, clamp=False)
         assert output_signal.max() > 1
+
+
+class Spectrogram(common_utils.TestBaseMixin):
+    @parameterized.expand([(0., ), (1., ), (2., ), (3., )])
+    def test_grad_at_zero(self, power):
+        """The gradient of power spectrogram should not be nan but zero near x=0
+
+        https://github.com/pytorch/audio/issues/993
+        """
+        x = torch.zeros(1, 22050, requires_grad=True)
+        spec = F.spectrogram(
+            x,
+            pad=0,
+            window=None,
+            n_fft=2048,
+            hop_length=None,
+            win_length=None,
+            power=power,
+            normalized=False,
+        )
+        spec.sum().backward()
+        assert not x.grad.isnan().sum()