Fix buffer size of DynamicBucketingSampler (#1468)

* Fix buffer size

* Fix for flake8

---------

Co-authored-by: yifanyeung <yifanyeung@yifanyeung.local>

egs/aidatatang_200zh/ASR/pruned_transducer_stateless2/asr_datamodule.py

@@ -288,8 +288,9 @@ class Aidatatang_200zhAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
-                buffer_size=50000,
             )
         else:
             logging.info("Using SimpleCutSampler.")

egs/aishell/ASR/tdnn_lstm_ctc/asr_datamodule.py

@@ -275,6 +275,8 @@ class AishellAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/aishell/ASR/transducer_stateless_modified-2/asr_datamodule.py

@@ -226,6 +226,8 @@ class AsrDataModule:
             max_duration=self.args.max_duration,
             shuffle=self.args.shuffle,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=True,
         )
 

egs/aishell2/ASR/pruned_transducer_stateless5/asr_datamodule.py

@@ -296,6 +296,8 @@ class AiShell2AsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/aishell4/ASR/pruned_transducer_stateless5/asr_datamodule.py

@@ -306,7 +306,8 @@ class Aishell4AsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
-                buffer_size=100000,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/alimeeting/ASR/pruned_transducer_stateless2/asr_datamodule.py

@@ -288,7 +288,8 @@ class AlimeetingAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
-                buffer_size=30000,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py

@@ -263,6 +263,8 @@ class AlimeetingAsrDataModule:
             max_cuts=self.args.max_cuts,
             shuffle=False,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=True,
         )
         logging.info("About to create train dataloader")

egs/ami/ASR/pruned_transducer_stateless7/asr_datamodule.py

@@ -269,6 +269,8 @@ class AmiAsrDataModule:
             max_cuts=self.args.max_cuts,
             shuffle=False,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=True,
         )
         logging.info("About to create train dataloader")

egs/ami/SURT/dprnn_zipformer/asr_datamodule.py

@@ -254,6 +254,8 @@ class AmiAsrDataModule:
                 max_cuts=self.args.max_cuts,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/commonvoice/ASR/pruned_transducer_stateless7/asr_datamodule.py

@@ -308,6 +308,8 @@ class CommonVoiceAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/commonvoice/ASR/pruned_transducer_stateless7_streaming/commonvoice_fr.py

@@ -310,6 +310,8 @@ class CommonVoiceAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/csj/ASR/local/utils/asr_datamodule.py

@@ -336,6 +336,8 @@ class CSJAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py

@@ -261,6 +261,8 @@ class GigaSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py

@@ -294,6 +294,8 @@ class GigaSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/gigaspeech/ASR/zipformer/asr_datamodule.py

@@ -311,6 +311,8 @@ class GigaSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/libricss/SURT/dprnn_zipformer/asr_datamodule.py

@@ -256,6 +256,8 @@ class LibriCssAsrDataModule:
                 max_cuts=self.args.max_cuts,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/libriheavy/ASR/zipformer/asr_datamodule.py

@@ -310,6 +310,8 @@ class LibriHeavyAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/libriheavy/ASR/zipformer_prompt_asr/asr_datamodule.py

@@ -341,6 +341,8 @@ class LibriHeavyAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/librispeech/ASR/pruned2_knowledge/asr_datamodule.py

@@ -286,6 +286,8 @@ class LibriSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 bucket_method="equal_duration",
                 drop_last=True,
             )

egs/librispeech/ASR/pruned_transducer_stateless3/asr_datamodule.py

@@ -223,6 +223,8 @@ class AsrDataModule:
             max_duration=self.args.max_duration,
             shuffle=self.args.shuffle,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=True,
         )
 
@@ -256,6 +258,8 @@ class AsrDataModule:
             max_duration=self.args.max_duration,
             shuffle=False,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=False,
         )
         logging.info("About to create dev dataloader")
@@ -282,6 +286,8 @@ class AsrDataModule:
             max_duration=self.args.max_duration,
             shuffle=False,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
         )
         logging.debug("About to create test dataloader")
         test_dl = DataLoader(

egs/librispeech/ASR/pruned_transducer_stateless7/gigaspeech.py

@@ -294,6 +294,8 @@ class GigaSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/librispeech/ASR/tdnn_lstm_ctc/asr_datamodule.py

@@ -311,6 +311,8 @@ class LibriSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/librispeech/ASR/tiny_transducer_ctc/asr_datamodule.py

@@ -304,6 +304,8 @@ class LibriSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/librispeech/WSASR/conformer_ctc2/asr_datamodule.py

@@ -227,6 +227,8 @@ class LibriSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/ljspeech/TTS/vits/tts_datamodule.py

@@ -196,6 +196,8 @@ class LJSpeechTtsDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/mgb2/ASR/conformer_ctc/asr_datamodule.py

@@ -266,6 +266,8 @@ class MGB2AsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/multi_zh-hans/ASR/zipformer/asr_datamodule.py

@@ -297,6 +297,8 @@ class AsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/multi_zh_en/ASR/zipformer/asr_datamodule.py

@@ -294,6 +294,8 @@ class AsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/spgispeech/ASR/pruned_transducer_stateless2/asr_datamodule.py

@@ -236,6 +236,8 @@ class SPGISpeechAsrDataModule:
             max_duration=self.args.max_duration,
             shuffle=False,
             num_buckets=self.args.num_buckets,
+            buffer_size=self.args.num_buckets * 2000,
+            shuffle_buffer_size=self.args.num_buckets * 5000,
             drop_last=True,
         )
         logging.info("About to create train dataloader")

egs/swbd/ASR/conformer_ctc/asr_datamodule.py

@@ -298,8 +298,9 @@ class SwitchBoardAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
-                buffer_size=50000,
             )
         else:
             logging.info("Using SimpleCutSampler.")

egs/tal_csasr/ASR/pruned_transducer_stateless5/asr_datamodule.py

@@ -306,8 +306,9 @@ class TAL_CSASRAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 num_cuts_for_bins_estimate=20000,
-                buffer_size=60000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/tedlium3/ASR/transducer_stateless/asr_datamodule.py

@@ -256,6 +256,8 @@ class TedLiumAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/timit/ASR/tdnn_lstm_ctc/asr_datamodule.py

@@ -222,6 +222,8 @@ class TimitAsrDataModule(DataModule):
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/vctk/TTS/vits/tts_datamodule.py

@@ -204,6 +204,8 @@ class VctkTtsDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/wenetspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py

@@ -292,7 +292,8 @@ class WenetSpeechAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
-                buffer_size=300000,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else:

egs/xbmu_amdo31/ASR/pruned_transducer_stateless5/asr_datamodule.py

@@ -296,6 +296,8 @@ class Xbmu_AmdoAsrDataModule:
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=self.args.drop_last,
             )
         else:

egs/yesno/ASR/tdnn/asr_datamodule.py

@@ -193,6 +193,8 @@ class YesNoAsrDataModule(DataModule):
                 max_duration=self.args.max_duration,
                 shuffle=self.args.shuffle,
                 num_buckets=self.args.num_buckets,
+                buffer_size=self.args.num_buckets * 2000,
+                shuffle_buffer_size=self.args.num_buckets * 5000,
                 drop_last=True,
             )
         else: