From 05ea517c1fba25e7dd3e29e50ec65adaed2c9410 Mon Sep 17 00:00:00 2001
From: PingFeng Luo <pingfengluo@gmail.com>
Date: Sun, 2 Jan 2022 13:12:42 +0800
Subject: [PATCH] try on-the-fly feature extraction

---
 egs/wenetspeech/ASR/tdnn_lstm_ctc/asr_datamodule.py | 13 ++++++++++++-
 egs/wenetspeech/ASR/transducer_stateless/train.py   |  7 +++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/egs/wenetspeech/ASR/tdnn_lstm_ctc/asr_datamodule.py b/egs/wenetspeech/ASR/tdnn_lstm_ctc/asr_datamodule.py
index 4de55b32e..eb581a647 100644
--- a/egs/wenetspeech/ASR/tdnn_lstm_ctc/asr_datamodule.py
+++ b/egs/wenetspeech/ASR/tdnn_lstm_ctc/asr_datamodule.py
@@ -17,11 +17,18 @@
 
 import argparse
 import logging
+import torch
 from functools import lru_cache
 from pathlib import Path
 from typing import List
 
-from lhotse import CutSet, Fbank, FbankConfig, load_manifest
+from lhotse import (
+    CutSet,
+    Fbank,
+    FbankConfig,
+    load_manifest,
+    set_caching_enabled,
+)
 from lhotse.dataset import (
     DynamicBucketingSampler,
     CutConcatenate,
@@ -37,6 +44,10 @@ from torch.utils.data import DataLoader
 from icefall.utils import str2bool
 
 
+set_caching_enabled(False)
+torch.set_num_threads(1)
+
+
 class WenetSpeechDataModule:
     """
     DataModule for k2 ASR experiments.
diff --git a/egs/wenetspeech/ASR/transducer_stateless/train.py b/egs/wenetspeech/ASR/transducer_stateless/train.py
index 9eb6deff4..ddde9b897 100755
--- a/egs/wenetspeech/ASR/transducer_stateless/train.py
+++ b/egs/wenetspeech/ASR/transducer_stateless/train.py
@@ -129,6 +129,13 @@ def get_parser():
         "2 means tri-gram",
     )
 
+    parser.add_argument(
+        "--on-the-fly",
+        type=str2bool,
+        default=True,
+        help="Use on-the-fly feature extraction",
+    )
+
     return parser