Repository: pytorch/audio

Issue:
torchaudio.transforms.TimeMasking is inplace
## 🐛 Bug

The function `torchaudio.transforms.TimeMasking` transforms the Tensor inline without expectation.

## To Reproduce

Steps to reproduce the behavior:

```python
import matplotlib.pyplot as plt
import librosa
import torchaudio
import torch
fmask = torchaudio.transforms.TimeMasking(time_mask_param=5)

BN = 2
max_frames = 200
data = torch.rand((BN, max_frames*160+240))
mt = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160,
                                                 window_fn=torch.hamming_window, n_mels=64)
spec = mt(data)

def plot_spectrogram(spec, title=None, ylabel='freq_bin', aspect='auto', xmax=None):
    fig, axs = plt.subplots(1, 1)
    axs.set_title(title or 'Spectrogram (db)')
    axs.set_ylabel(ylabel)
    axs.set_xlabel('frame')
    im = axs.imshow(librosa.power_to_db(spec), origin='lower', aspect=aspect)
    if xmax:
        axs.set_xlim((0, xmax))
        fig.colorbar(im, ax=axs)
        plt.show(block=False)

spec = mt(data)

fmask(spec) # expect `mask_spec = fmask(spec)` or `fmask(spec, inline=False)`
print((spec==mt(data)).all()) # False (but  True is expected)
```

## Expected behavior

expect `mask_spec = fmask(spec)` or `fmask(spec, inline=False)`

## Environment
```
Collecting environment information...
PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: version 3.5.1

Python version: 3.8 (64-bit runtime)
Is CUDA available: True

Versions of relevant libraries:
[pip3] numpy==1.20.2
[pip3] pytorch-metric-learning==0.9.97
[pip3] torch==1.8.1
[pip3] torch-scatter==2.0.6
[pip3] torchaudio==0.8.1
[pip3] torchlars==0.1.2
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl    defaults
[conda] cudatoolkit               10.1.243             h6bb024c_0    defaults
[conda] mkl                       2020.2                      256    defaults
[conda] mkl-service               2.3.0            py38he904b0f_0    defaults
[conda] mkl_fft                   1.2.0            py38h23d657b_0    defaults
[conda] mkl_random                1.1.1            py38h0573a6f_0    defaults
[conda] numpy                     1.19.2           py38h54aff64_0    defaults
[conda] numpy-base                1.19.2           py38hfa32c7d_0    defaults
[conda] pytorch                   1.6.0           py3.8_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] pytorch-metric-learning   0.9.97                   pypi_0    pypi
[conda] torch-scatter             2.0.6                    pypi_0    pypi
[conda] torchaudio                0.6.0                      py38    pytorch
[conda] torchlars                 0.1.2                    pypi_0    pypi
[conda] torchvision               0.7.0                py38_cu101    pytorch
```




Discussion:
do you mean inplace?
following up on @gchanan's question, for a functional form of the transform, we also have [torchaudio.functional.mask_along_axis](https://pytorch.org/audio/stable/functional.html#mask-along-axis) and [torchaudio.functional.mask_along_axis_iid](https://pytorch.org/audio/stable/functional.html#mask-along-axis-iid). this is the functional that does the masking in the transform too, see [here](https://github.com/pytorch/audio/blob/master/torchaudio/transforms.py#L857). is that what you were looking for @seacj ?
Hi @seacj 

Thanks for the report. I am observing something interesting here. When I run your snippet, I do get `False` for the last statement, which is not expected, but when I simplify the script by removing the `matplotlib` and `librosa` it says `True`.

```python
import torchaudio
import torch

fmask = torchaudio.transforms.TimeMasking(time_mask_param=5)

BN = 2
max_frames = 200
data = torch.rand((BN, max_frames * 160 + 240))
mt = torchaudio.transforms.MelSpectrogram(
    sample_rate=16000,
    n_fft=512,
    win_length=400,
    hop_length=160,
    window_fn=torch.hamming_window,
    n_mels=64,
)
spec = mt(data)

fmask(spec)
print((spec == mt(data)).all())
```


```
$ python foo.py
tensor(True)
```
@seacj I could simplify the script, and indeed I see un-deterministic behavoir

```python
import torchaudio
import torch

spec = torch.rand((2, 64, 102))
fmask = torchaudio.transforms.TimeMasking(time_mask_param=5)
spec_ref = spec.clone()
fmask(spec)
torch.testing.assert_allclose(spec, spec_ref)
print('success')
```

```
$ for _ in `seq 1 10`; do python foo.py; done
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 256 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9962434768676758 (0.0 vs. 0.9962434768676758), which occurred at index (0, 56, 8).
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 512 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9975877404212952 (0.0 vs. 0.9975877404212952), which occurred at index (0, 50, 62).
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 512 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9955663084983826 (0.0 vs. 0.9955663084983826), which occurred at index (0, 6, 53).
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 256 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.99956214427948 (0.0 vs. 0.99956214427948), which occurred at index (1, 4, 1).
success
success
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 256 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9982173442840576 (0.0 vs. 0.9982173442840576), which occurred at index (0, 54, 96).
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 512 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9988617897033691 (0.0 vs. 0.9988617897033691), which occurred at index (1, 27, 17).
success
Traceback (most recent call last):
  File "foo.py", line 8, in <module>
    torch.testing.assert_allclose(spec, spec_ref)
  File "/fsx/users/moto/conda/lib/python3.8/site-packages/torch/testing/_core.py", line 270, in assert_allclose
    raise AssertionError(msg)
AssertionError: With rtol=0.0001 and atol=1e-05, found 512 element(s) (out of 13056) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 0.9927456378936768 (0.0 vs. 0.9927456378936768), which occurred at index (0, 0, 81).
```
I do see there is a in-place operation.

https://github.com/pytorch/audio/blob/b5d802795ffb381947e9fbbc1b1da40aa6c64b5f/torchaudio/functional/functional.py#L748-L750
This is also changing the tensor in place

https://github.com/pytorch/audio/blob/b5d802795ffb381947e9fbbc1b1da40aa6c64b5f/torchaudio/functional/functional.py#L790
@mthrok 
> I could simplify the script, and indeed I see un-deterministic behavior

The nondeterministic behavior is caused by the randomness in the algorithm. In your example, `spec` will always be modified in-place as long as there is any masking. The `success` comes from cases in which no columns end up being masked (when `mask_start` equals `mask_end`), and therefore `spec` stays the same.

In many cases, `reshape` will not return a copy of the tensor, so modifying `specgram` later on will change the data in the original tensor as well. We should modify how masking is being done so that it doesn't modify the original tensor.

https://github.com/pytorch/audio/blob/b5d802795ffb381947e9fbbc1b1da40aa6c64b5f/torchaudio/functional/functional.py#L778


> @mthrok
> 
> > I could simplify the script, and indeed I see un-deterministic behavior
> 
> The nondeterministic behavior is caused by the randomness in the algorithm. In your example, `spec` will always be modified in-place as long as there is any masking. The `success` comes from cases in which no columns end up being masked (when `mask_start` equals `mask_end`), and therefore `spec` stays the same.
> 
> In many cases, `reshape` will not return a copy of the tensor, so modifying `specgram` later on will change the data in the original tensor as well. We should modify how masking is being done so that it doesn't modify the original tensor.
> 
> https://github.com/pytorch/audio/blob/b5d802795ffb381947e9fbbc1b1da40aa6c64b5f/torchaudio/functional/functional.py#L778

That‘s right. So `plot_spectrogram(spec[0])` helps see what happens. 

Resolving patch:
diff --git a/torchaudio/functional/functional.py b/torchaudio/functional/functional.py
--- a/torchaudio/functional/functional.py
+++ b/torchaudio/functional/functional.py
@@ -746,7 +746,7 @@ def mask_along_axis_iid(
 
     # Per batch example masking
     specgrams = specgrams.transpose(axis, -1)
-    specgrams.masked_fill_((mask >= mask_start) & (mask < mask_end), mask_value)
+    specgrams = specgrams.masked_fill((mask >= mask_start) & (mask < mask_end), mask_value)
     specgrams = specgrams.transpose(axis, -1)
 
     return specgrams
@@ -772,24 +772,25 @@ def mask_along_axis(
     Returns:
         Tensor: Masked spectrogram of dimensions (channel, freq, time)
     """
+    if axis != 1 and axis != 2:
+        raise ValueError('Only Frequency and Time masking are supported')
 
     # pack batch
     shape = specgram.size()
     specgram = specgram.reshape([-1] + list(shape[-2:]))
-
     value = torch.rand(1) * mask_param
     min_value = torch.rand(1) * (specgram.size(axis) - value)
 
     mask_start = (min_value.long()).squeeze()
     mask_end = (min_value.long() + value.long()).squeeze()
+    mask = torch.arange(0, specgram.shape[axis], device=specgram.device, dtype=specgram.dtype)
+    mask = (mask >= mask_start) & (mask < mask_end)
+    if axis == 1:
+        mask = mask.unsqueeze(-1)
 
     assert mask_end - mask_start < mask_param
-    if axis == 1:
-        specgram[:, mask_start:mask_end] = mask_value
-    elif axis == 2:
-        specgram[:, :, mask_start:mask_end] = mask_value
-    else:
-        raise ValueError('Only Frequency and Time masking are supported')
+
+    specgram = specgram.masked_fill(mask, mask_value)
 
     # unpack batch
     specgram = specgram.reshape(shape[:-2] + specgram.shape[-2:])


Test patch:
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
@@ -227,6 +227,38 @@ def test_mask_along_axis_iid(self, mask_param, mask_value, axis):
         assert mask_specgrams.size() == specgrams.size()
         assert (num_masked_columns < mask_param).sum() == num_masked_columns.numel()
 
+    @parameterized.expand(
+        list(itertools.product([(2, 1025, 400), (1, 201, 100)], [100], [0., 30.], [1, 2]))
+    )
+    def test_mask_along_axis_preserve(self, shape, mask_param, mask_value, axis):
+        """mask_along_axis should not alter original input Tensor
+
+        Test is run 5 times to bound the probability of no masking occurring to 1e-10
+        See https://github.com/pytorch/audio/issues/1478
+        """
+        torch.random.manual_seed(42)
+        for _ in range(5):
+            specgram = torch.randn(*shape, dtype=self.dtype, device=self.device)
+            specgram_copy = specgram.clone()
+            F.mask_along_axis(specgram, mask_param, mask_value, axis)
+
+            self.assertEqual(specgram, specgram_copy)
+
+    @parameterized.expand(list(itertools.product([100], [0., 30.], [2, 3])))
+    def test_mask_along_axis_iid_preserve(self, mask_param, mask_value, axis):
+        """mask_along_axis_iid should not alter original input Tensor
+
+        Test is run 5 times to bound the probability of no masking occurring to 1e-10
+        See https://github.com/pytorch/audio/issues/1478
+        """
+        torch.random.manual_seed(42)
+        for _ in range(5):
+            specgrams = torch.randn(4, 2, 1025, 400, dtype=self.dtype, device=self.device)
+            specgrams_copy = specgrams.clone()
+            F.mask_along_axis_iid(specgrams, mask_param, mask_value, axis)
+
+            self.assertEqual(specgrams, specgrams_copy)
+
 
 class FunctionalComplex(TestBaseMixin):
     complex_dtype = None
diff --git a/test/torchaudio_unittest/functional/torchscript_consistency_impl.py b/test/torchaudio_unittest/functional/torchscript_consistency_impl.py
--- a/test/torchaudio_unittest/functional/torchscript_consistency_impl.py
+++ b/test/torchaudio_unittest/functional/torchscript_consistency_impl.py
@@ -13,7 +13,7 @@
 
 
 class Functional(TempDirMixin, TestBaseMixin):
-    """Implements test for `functinoal` modul that are performed for different devices"""
+    """Implements test for `functional` module that are performed for different devices"""
     def _assert_consistency(self, func, tensor, shape_only=False):
         tensor = tensor.to(device=self.device, dtype=self.dtype)
 
@@ -21,8 +21,12 @@ def _assert_consistency(self, func, tensor, shape_only=False):
         torch.jit.script(func).save(path)
         ts_func = torch.jit.load(path)
 
+        torch.random.manual_seed(40)
         output = func(tensor)
+
+        torch.random.manual_seed(40)
         ts_output = ts_func(tensor)
+
         if shape_only:
             ts_output = ts_output.shape
             output = output.shape