More fixes to gigaspeech recipe

Author: pkufool

egs/gigaspeech/ASR/zipformer/train.py

@@ -416,6 +416,17 @@ def get_parser():
         help="Accumulate stats on activations, print them and exit.",
     )
 
+    parser.add_argument(
+        "--scan-for-oom-batches",
+        type=str2bool,
+        default=False,
+        help="""
+        Whether to scan for oom batches before training, this is helpful for
+        finding the suitable max_duration, you only need to run it once.
+        Caution: a little time consuming.
+        """,
+    )
+
     parser.add_argument(
         "--inf-check",
         type=str2bool,
@@ -1197,14 +1208,14 @@ def remove_short_utt(c: Cut):
     valid_cuts = valid_cuts.filter(remove_short_utt)
     valid_dl = gigaspeech.valid_dataloaders(valid_cuts)
 
-    # if not params.print_diagnostics:
-    #    scan_pessimistic_batches_for_oom(
-    #        model=model,
-    #        train_dl=train_dl,
-    #        optimizer=optimizer,
-    #        sp=sp,
-    #        params=params,
-    #    )
+    if not params.print_diagnostics and params.scan_for_oom_batches:
+        scan_pessimistic_batches_for_oom(
+            model=model,
+            train_dl=train_dl,
+            optimizer=optimizer,
+            sp=sp,
+            params=params,
+        )
 
     scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
     if checkpoints and "grad_scaler" in checkpoints:

egs/gigaspeech/KWS/zipformer/asr_datamodule.py

@@ -1,5 +1,5 @@
 # Copyright      2021  Piotr Żelasko
-# Copyright      2023  Xiaomi Corporation     (Author: Yifan Yang)
+# Copyright      2024  Xiaomi Corporation     (Author: Wei Kang)
 #
 # See ../../../../LICENSE for clarification regarding multiple authors
 #
@@ -448,13 +448,6 @@ def test_cuts(self) -> CutSet:
             self.args.manifest_dir / "gigaspeech_cuts_TEST.jsonl.gz"
         )
 
-    @lru_cache()
-    def libri_100_cuts(self) -> CutSet:
-        logging.info("About to get libri100 cuts")
-        return load_manifest_lazy(
-            self.args.manifest_dir / "librispeech_cuts_train-clean-100.jsonl.gz"
-        )
-
     @lru_cache()
     def fsc_train_cuts(self) -> CutSet:
         logging.info("About to get fluent speech commands train cuts")

egs/gigaspeech/KWS/zipformer/decode.py

@@ -274,7 +274,7 @@ def decode_one_batch(
         model=model,
         encoder_out=encoder_out,
         encoder_out_lens=encoder_out_lens,
-        context_graph=kws_graph,
+        keywords_graph=kws_graph,
         beam=params.beam,
         num_tailing_blanks=params.num_tailing_blanks,
         blank_penalty=params.blank_penalty,

egs/gigaspeech/KWS/zipformer/decode-asr.py egs/gigaspeech/KWS/zipformer/decode_asr.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 #
-# Copyright 2021-2023 Xiaomi Corporation (Author: Fangjun Kuang,
-#                                                 Zengwei Yao)
+# Copyright 2021-2024 Xiaomi Corporation (Author: Fangjun Kuang,
+#                                                 Zengwei Yao,
+#                                                 Wei Kang)
 #
 # See ../../../../LICENSE for clarification regarding multiple authors
 #