remove batch for mask.

Author: vincentqb

test/test_functional.py

@@ -560,28 +560,6 @@ def _test_batch(self, functional, tensor, *args, **kwargs):
         torch.random.manual_seed(42)
         computed = functional(tensors.clone(), *args, **kwargs)
 
-        self._compare_estimate(computed, expected, **kwargs_compare)
-
-    def test_batch_mask_along_axis_iid(self):
-
-        mask_param = 2
-        mask_value = 30.
-        axis = 2
-
-        tensor = torch.rand(2, 5, 5)
-
-        self._test_batch(F.mask_along_axis_iid, tensor, mask_param=mask_param, mask_value=mask_value, axis=axis, atol=1e-1, rtol=1e-1)
-
-    def test_batch_mask_along_axis(self):
-
-        tensor = torch.rand(2, 5, 5)
-
-        mask_param = 2
-        mask_value = 30.
-        axis = 2
-
-        self._test_batch(F.mask_along_axis, tensor, mask_param=mask_param, mask_value=mask_value, axis=axis)
-
     def test_torchscript_create_fb_matrix(self):
 
         n_stft = 100

torchaudio/functional.py

@@ -835,21 +835,14 @@ def mask_along_axis_iid(specgrams, mask_param, mask_value, axis):
         axis (int): Axis to apply masking on (2 -> frequency, 3 -> time)
 
     Returns:
-        torch.Tensor: Masked spectrograms of dimensions (..., channel, freq, time)
+        torch.Tensor: Masked spectrograms of dimensions (batch, channel, freq, time)
     """
 
-    # pack batch
-    shape = specgrams.size()
-    specgrams = specgrams.reshape([-1] + list(shape[-3:]))
-
     if axis != 2 and axis != 3:
         raise ValueError('Only Frequency and Time masking are supported')
 
-    # Shift so as to start from the end
-    axis -= 4
-
-    value = torch.rand(specgrams.shape[:-2]) * mask_param
-    min_value = torch.rand(specgrams.shape[:-2]) * (specgrams.size(axis) - value)
+    value = torch.rand(specgrams.shape[:2]) * mask_param
+    min_value = torch.rand(specgrams.shape[:2]) * (specgrams.size(axis) - value)
 
     # Create broadcastable mask
     mask_start = (min_value.long())[..., None, None].float()
@@ -861,9 +854,6 @@ def mask_along_axis_iid(specgrams, mask_param, mask_value, axis):
     specgrams.masked_fill_((mask >= mask_start) & (mask < mask_end), mask_value)
     specgrams = specgrams.transpose(axis, -1)
 
-    # unpack batch
-    specgrams = specgrams.reshape(shape[:-3] + specgrams.shape[-3:])
-
     return specgrams
 
 
@@ -875,19 +865,15 @@ def mask_along_axis(specgram, mask_param, mask_value, axis):
     All examples will have the same mask interval.
 
     Args:
-        specgram (Tensor): Real spectrogram (..., channel, freq, time)
+        specgram (Tensor): Real spectrogram (channel, freq, time)
         mask_param (int): Number of columns to be masked will be uniformly sampled from [0, mask_param]
         mask_value (float): Value to assign to the masked columns
         axis (int): Axis to apply masking on (1 -> frequency, 2 -> time)
 
     Returns:
-        torch.Tensor: Masked spectrogram of dimensions (..., channel, freq, time)
+        torch.Tensor: Masked spectrogram of dimensions (channel, freq, time)
     """
 
-    # pack batch
-    shape = specgram.size()
-    specgram = specgram.reshape([-1] + list(shape[-3:]))
-
     value = torch.rand(1) * mask_param
     min_value = torch.rand(1) * (specgram.size(axis) - value)
 
@@ -902,9 +888,6 @@ def mask_along_axis(specgram, mask_param, mask_value, axis):
     else:
         raise ValueError('Only Frequency and Time masking are supported')
 
-    # unpack batch
-    specgram = specgram.reshape(shape[:-3] + specgram.shape[-3:])
-
     return specgram