update

examples/vm_sound_classification/conv2d_classifier.yaml

@@ -10,6 +10,11 @@ mel_spectrogram_param:
   window_fn: hamming
   n_mels: 80
 
+spec_augment_param:
+  aug_volume_factor_range:
+    - 0.5
+    - 2.0
+
 conv2d_block_param_list:
   - batch_norm: true
     in_channels: 1

examples/vm_sound_classification/step_1_prepare_data.py

@@ -70,7 +70,7 @@ def get_dataset(args):
         "mute": "mute",
         "noise": "noise",
         "noise_mute": "noise",
-        "voice": "voice_or_noise",
+        "voice": "voice",
         "voicemail": "voicemail",
     }
     # label8_map = {

toolbox/torchaudio/augment/spec_augment.py

@@ -0,0 +1,44 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/wenet-e2e/wenet/blob/main/wenet/dataset/processor.py
+"""
+import random
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class SpecAugment(nn.Module):
+    def __init__(self,
+                 aug_volume_factor_range: Tuple[float, float] = (0.5, 2.0),
+                 ):
+        super().__init__()
+        self.aug_volume_factor_range = aug_volume_factor_range
+
+    @staticmethod
+    def augment_volume(spec: torch.Tensor, factor_range: Tuple[float, float] = (0.5, 2.0)):
+        factor = random.uniform(*factor_range)
+        spec_ = spec.clone().detach()
+        spec_ *= factor
+        return spec_
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        spec = self.augment_volume(spec, self.aug_volume_factor_range)
+        return spec
+
+
+def main():
+    spec_augment = SpecAugment()
+
+    spec = torch.randn(size=(1, 10, 4))
+    print(spec)
+
+    spec_ = spec_augment.forward(spec)
+    print(spec_)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/cnn_audio_classifier/configuration_cnn_audio_classifier.py

@@ -8,6 +8,7 @@ from toolbox.torchaudio.configuration_utils import PretrainedConfig
 class CnnAudioClassifierConfig(PretrainedConfig):
     def __init__(self,
                  mel_spectrogram_param: dict,
+                 spec_augment_param: dict,
                  cls_head_param: dict,
                  conv1d_block_param_list: List[dict] = None,
                  conv2d_block_param_list: List[dict] = None,
@@ -15,6 +16,7 @@ class CnnAudioClassifierConfig(PretrainedConfig):
                  ):
         super(CnnAudioClassifierConfig, self).__init__(**kwargs)
         self.mel_spectrogram_param = mel_spectrogram_param
+        self.spec_augment_param = spec_augment_param
         self.cls_head_param = cls_head_param
         self.conv1d_block_param_list = conv1d_block_param_list
         self.conv2d_block_param_list = conv2d_block_param_list

toolbox/torchaudio/models/cnn_audio_classifier/modeling_cnn_audio_classifier.py

@@ -9,6 +9,7 @@ import torchaudio
 
 from toolbox.torchaudio.models.cnn_audio_classifier.configuration_cnn_audio_classifier import CnnAudioClassifierConfig
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.augment.spec_augment import SpecAugment
 
 
 MODEL_FILE = "model.pt"
@@ -240,6 +241,7 @@ class SpectrogramEncoder(nn.Module):
 class WaveEncoder(nn.Module):
     def __init__(self,
                  mel_spectrogram_param: dict,
+                 spec_augment_param: dict,
                  conv1d_block_param_list: List[dict] = None,
                  conv2d_block_param_list: List[dict] = None,
                  ):
@@ -262,6 +264,10 @@ class WaveEncoder(nn.Module):
             ),
         )
 
+        self.spec_augment = SpecAugment(
+            aug_volume_factor_range=spec_augment_param["aug_volume_factor_range"]
+        )
+
         self.spectrogram_encoder = SpectrogramEncoder(
             conv1d_block_param_list=conv1d_block_param_list,
             conv2d_block_param_list=conv2d_block_param_list,
@@ -277,6 +283,9 @@ class WaveEncoder(nn.Module):
             x = x.log()
             x = x - torch.mean(x, dim=-1, keepdim=True)
 
+            if self.training:
+                x = self.spec_augment.forward(x)
+
         x = x.transpose(1, 2)
 
         features = self.spectrogram_encoder.forward(x)
@@ -346,6 +355,7 @@ class WaveClassifierPretrainedModel(WaveClassifier):
         super(WaveClassifierPretrainedModel, self).__init__(
             wave_encoder=WaveEncoder(
                 mel_spectrogram_param=config.mel_spectrogram_param,
+                spec_augment_param=config.spec_augment_param,
                 conv1d_block_param_list=config.conv1d_block_param_list,
                 conv2d_block_param_list=config.conv2d_block_param_list,
             ),