Feature/adaptive temporal masks (#7)

* introduce adaptive_temporal_masks to do mask depending on speech length

* bump gx version

Author: sutetako

nemo/collections/asr/parts/submodules/spectr_augment.py

@@ -83,6 +83,14 @@ def __init__(
         self.mask_value = mask_value
         self.use_vectorized_code = use_vectorized_code
 
+        if isinstance(time_masks, int):
+            self.adaptive_temporal_masks = False
+        else:
+            if time_masks > 1.0 or time_masks < 0.0:
+                raise ValueError("If `time_masks` is a float value, must be in range [0, 1]")
+
+            self.adaptive_temporal_masks = True
+
         if isinstance(time_width, int):
             self.adaptive_temporal_width = False
         else:
@@ -114,6 +122,12 @@ def _forward_legacy(self, input_spec, length):
                 width = self._rng.randint(0, self.freq_width)
                 fill_mask[idx, start : start + width, :] = True
 
+            # Derive the number of masks, sometimes based percentage of input length.
+            if self.adaptive_temporal_masks:
+                time_max_masks = int(lengths_cpu[idx] * self.time_masks)
+            else:
+                time_max_masks = self.time_masks
+
             # Derive time width, sometimes based percentage of input length.
             if self.adaptive_temporal_width:
                 time_max_width = max(1, int(lengths_cpu[idx] * self.time_width))
@@ -122,7 +136,7 @@ def _forward_legacy(self, input_spec, length):
             time_start_upper_bound = max(1, lengths_cpu[idx] - time_max_width)
 
             # Set time masking
-            for _ in range(self.time_masks):
+            for _ in range(time_max_masks):
                 start = self._rng.randint(0, time_start_upper_bound)
                 width = self._rng.randint(0, time_max_width)
                 fill_mask[idx, :, start : start + width] = True

nemo/package_info.py

@@ -16,7 +16,7 @@
 MAJOR = 2
 MINOR = 2
 PATCH = 0
-PRE_RELEASE = ''
+PRE_RELEASE = '+gx1'
 
 # Use the following formatting: (major, minor, patch, pre-release)
 VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)