add dropout to linear decoder, clean up

espnet · Shikhar-S · Dec 10, 2024 · Dec 10, 2024 · Dec 10, 2024 · Dec 10, 2024
commit ca9ca1ef8bc86753238ca7a0de05d87b8f57abb3
diff --git a/egs2/esc50/asr1/conf/beats_classification.yaml b/egs2/esc50/asr1/conf/beats_classification.yaml
@@ -2,33 +2,32 @@ token_type: word
 
 optim: adamw
 optim_conf:
-    lr: 3.0e-5
+    lr: 1.0e-4
     weight_decay: 1.0e-2
     betas: [0.9, 0.98]
 
 accum_grad: 1
 
-batch_size: 60 # 300 sec (5 sec per sample) # 1600 samples in train set ==> 1600/60 = 27 steps per epoch
-max_epoch: 2962 # originally 80K steps, @27 steps per epoch ~2962
-
+batch_size: 128 # 12.5 steps per epoch with 1600 samples
+max_epoch: 1000
 
 scheduler: CosineAnnealingWarmupRestarts
 scheduler_conf:
-    first_cycle_steps: 80000 # Only one cycle, so this is last epoch step
-    warmup_steps: 8000 # paper mentions epochs??!
-    max_lr: 3.0e-5
-    min_lr: 0
+    first_cycle_steps: 6000
+    warmup_steps: 300
+    max_lr: 1.0e-4
+    min_lr: 5.0e-6
 
 # BEATs implementation takes care of generating mel spectrogram, normalization and specaug
 frontend: none
 input_size: 1 # important to set input_size to 1 if frontend is none
 normalize: none # BEATs code does global mean and variance normalization
 
 # Initialization for the decoder
-init: normal
+init: xavier_normal
 
 model_conf:
-    ctc_weight: 0.0  # No CTC, only attention branch
+    ctc_weight: 0.0  # No CTC, no attention.
     lsm_weight: 0.1  # label smoothing weight
     length_normalized_loss: true
 
@@ -40,7 +39,7 @@ best_model_criterion:
 -   - valid
     - acc
     - max
-keep_nbest_models: 5
+keep_nbest_models: 1
 use_amp: false  # whether to use automatic mixed precision
 num_att_plot: 0
 num_workers: 2 # dataloader workers
@@ -50,28 +49,28 @@ encoder_conf:
     # Please download the BEATs model from https://github.com/microsoft/unilm/tree/master/beats
     # (iter3) and update the path below
     beats_ckpt_path: /compute/babel-13-33/sbharad2/models/BEATs/BEATs_iter3.pt
-    # values from Appendix A.1 of the BEATs paper
+    # Most values from Appendix A.1 of the BEATs paper or tuned on fold 5. 
+    # Please also check the README.md
     fbank_mean: 11.72215
     fbank_std: 10.60431
     beats_config:
-        layer_wise_gradient_decay_ratio: 0.2
+        layer_wise_gradient_decay_ratio: 0.2    
         encoder_layerdrop: 0.1
         dropout: 0.0
-    # TODO(shikhar): specaug = 0.2 ??, roll augmentation???
     specaug_config:
-        apply_freq_mask: false # true
-        freq_mask_width_range:
-        - 0
-        - 64
-        num_freq_mask: 2
+        apply_time_warp: true
+        apply_freq_mask: false
         apply_time_mask: true
         time_mask_width_ratio_range:
         - 0
-        - 0.12
-        num_time_mask: 5
+        - 0.06
+        num_time_mask: 1
+    roll_augment: true
+    roll_interval: 16000 # 1 second, only 5 possible augmentations per sample
     use_weighted_representation: false
 
 # Simple linear decoder for classification.
 decoder: linear_decoder
 decoder_conf:
     pooling: mean
+    dropout: 0.1
diff --git a/espnet2/asr/decoder/linear_decoder.py b/espnet2/asr/decoder/linear_decoder.py
@@ -20,13 +20,18 @@ def __init__(
         vocab_size: int,
         encoder_output_size: int,
         pooling: str = "CLS",
+        dropout: float = 0.0,
     ):
         """Initialize the module."""
         super().__init__()
 
         self.input_dim = encoder_output_size
         assert vocab_size > 3, "Invalid vocab size, must be > 3."
         self.output_dim = vocab_size - 3
+        self.dropout = None
+        if dropout != 0.0:
+            self.dropout = torch.nn.Dropout(p=dropout)
+        # self.output_dim = vocab_size
         self.linear_out = torch.nn.Linear(self.input_dim, self.output_dim)
         assert pooling in [
             "mean",
@@ -51,6 +56,8 @@ def forward(
         """
 
         mask = make_pad_mask(lengths=hlens, xs=hs_pad, length_dim=1).to(hs_pad.device)
+        if self.dropout is not None:
+            hs_pad = self.dropout(hs_pad)
         if self.pooling == "mean":
             unmasked_entries = (~mask).to(dtype=hs_pad.dtype)
             input_feature = (hs_pad * unmasked_entries).sum(dim=1)

diff --git a/espnet2/asr/encoder/beats_encoder.py b/espnet2/asr/encoder/beats_encoder.py
@@ -133,6 +133,9 @@ class BeatsEncoder(AbsEncoder):
         add_positional_information: Add learned positional embeddings.
         max_positions: Maximum number of positions for positional embeddings.
             Required if `add_positional_information` is True.
+        roll_augment: Apply roll augmentation to the input.
+        roll_interval: Interval for roll augmentation. All rolling is
+            quantized to this interval.
     """
 
     def __init__(
@@ -355,7 +358,6 @@ def forward(
             masks: None
         """
         if self.roll_augment and self.training:
-            # rolls quantized to 100 ms
             xs_pad = roll_tensor(
                 xs_pad.unsqueeze(-1), ilens, fixed_intervals=self.roll_interval
             ).squeeze(-1)