⚡ fix read audio, update dataset config and add transducer greedy v2

Author: nglehuy

examples/conformer/config.yml

@@ -24,12 +24,16 @@ speech_config:
   normalize_per_feature: False
 
 decoder_config:
-  vocabulary: null
-  target_vocab_size: 1024
+  vocabulary: ./vocabularies/librispeech_train_4_4076.subwords
+  target_vocab_size: 4096
   max_subword_length: 4
   blank_at_zero: True
   beam_width: 5
   norm_score: True
+  corpus_files:
+    - /media/nlhuy/Data/ML/ASR/Raw/LibriSpeech/LibriSpeech/train-clean-100/transcripts.tsv
+    - /media/nlhuy/Data/ML/ASR/Raw/LibriSpeech/LibriSpeech/train-clean-360/transcripts.tsv
+    - /media/nlhuy/Data/ML/ASR/Raw/LibriSpeech/LibriSpeech/train-other-500/transcripts.tsv
 
 model_config:
   name: conformer
@@ -53,32 +57,51 @@ model_config:
   prediction_rnn_units: 320
   prediction_rnn_type: lstm
   prediction_rnn_implementation: 2
-  prediction_layer_norm: True
+  prediction_layer_norm: False
   prediction_projection_units: 0
-  joint_dim: 320
+  joint_dim: 640
   joint_activation: tanh
 
 learning_config:
-  augmentations:
+  train_dataset_config:
     use_tf: True
-    after:
-      time_masking:
-        num_masks: 10
-        mask_factor: 100
-        p_upperbound: 0.05
-      freq_masking:
-        num_masks: 1
-        mask_factor: 27
-
-  dataset_config:
-    train_paths:
+    augmentation_config:
+      after:
+        time_masking:
+          num_masks: 10
+          mask_factor: 100
+          p_upperbound: 0.05
+        freq_masking:
+          num_masks: 1
+          mask_factor: 27
+    data_paths:
       - /mnt/Miscellanea/Datasets/Speech/LibriSpeech/train-clean-100/transcripts.tsv
-    eval_paths:
+    tfrecords_dir: /mnt/Miscellanea/Datasets/Speech/LibriSpeech/tfrecords-test
+    shuffle: True
+    cache: True
+    buffer_size: 100
+    drop_remainder: True
+
+  eval_dataset_config:
+    use_tf: True
+    data_paths:
       - /mnt/Miscellanea/Datasets/Speech/LibriSpeech/dev-clean/transcripts.tsv
       - /mnt/Miscellanea/Datasets/Speech/LibriSpeech/dev-other/transcripts.tsv
-    test_paths:
+    tfrecords_dir: /mnt/Miscellanea/Datasets/Speech/LibriSpeech/tfrecords-test
+    shuffle: False
+    cache: True
+    buffer_size: 100
+    drop_remainder: True
+
+  test_dataset_config:
+    use_tf: True
+    data_paths:
       - /mnt/Miscellanea/Datasets/Speech/LibriSpeech/test-clean/transcripts.tsv
     tfrecords_dir: /mnt/Miscellanea/Datasets/Speech/LibriSpeech/tfrecords-test
+    shuffle: False
+    cache: True
+    buffer_size: 100
+    drop_remainder: True
 
   optimizer_config:
     warmup_steps: 40000

examples/conformer/test_conformer.py

@@ -25,26 +25,19 @@
 
 parser = argparse.ArgumentParser(prog="Conformer Testing")
 
-parser.add_argument("--config", type=str, default=DEFAULT_YAML,
-                    help="The file path of model configuration file")
+parser.add_argument("--config", type=str, default=DEFAULT_YAML, help="The file path of model configuration file")
 
-parser.add_argument("--saved", type=str, default=None,
-                    help="Path to saved model")
+parser.add_argument("--saved", type=str, default=None, help="Path to saved model")
 
-parser.add_argument("--tfrecords", default=False, action="store_true",
-                    help="Whether to use tfrecords as dataset")
+parser.add_argument("--tfrecords", default=False, action="store_true", help="Whether to use tfrecords as dataset")
 
-parser.add_argument("--mxp", default=False, action="store_true",
-                    help="Enable mixed precision")
+parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
-parser.add_argument("--device", type=int, default=0,
-                    help="Device's id to run test on")
+parser.add_argument("--device", type=int, default=0, help="Device's id to run test on")
 
-parser.add_argument("--cpu", default=False, action="store_true",
-                    help="Whether to only use cpu")
+parser.add_argument("--cpu", default=False, action="store_true", help="Whether to only use cpu")
 
-parser.add_argument("--output_name", type=str, default="test",
-                    help="Result filename name prefix")
+parser.add_argument("--output_name", type=str, default="test", help="Result filename name prefix")
 
 args = parser.parse_args()
 
@@ -53,7 +46,7 @@
 setup_devices([args.device], cpu=args.cpu)
 
 from tensorflow_asr.configs.config import Config
-from tensorflow_asr.datasets.asr_dataset import ASRTFRecordTestDataset, ASRSliceTestDataset
+from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
 from tensorflow_asr.runners.base_runners import BaseTester
@@ -67,19 +60,14 @@
 assert args.saved
 
 if args.tfrecords:
-    test_dataset = ASRTFRecordTestDataset(
-        data_paths=config.learning_config.dataset_config.test_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="test", shuffle=False
+    test_dataset = ASRTFRecordDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.test_dataset_config)
     )
 else:
-    test_dataset = ASRSliceTestDataset(
-        data_paths=config.learning_config.dataset_config.test_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="test", shuffle=False
+    test_dataset = ASRSliceDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.test_dataset_config)
     )
 
 # build model

examples/conformer/test_subword_conformer.py

@@ -25,31 +25,23 @@
 
 parser = argparse.ArgumentParser(prog="Conformer Testing")
 
-parser.add_argument("--config", type=str, default=DEFAULT_YAML,
-                    help="The file path of model configuration file")
+parser.add_argument("--config", type=str, default=DEFAULT_YAML, help="The file path of model configuration file")
 
-parser.add_argument("--saved", type=str, default=None,
-                    help="Path to saved model")
+parser.add_argument("--saved", type=str, default=None, help="Path to saved model")
 
-parser.add_argument("--tfrecords", default=False, action="store_true",
-                    help="Whether to use tfrecords as dataset")
+parser.add_argument("--tfrecords", default=False, action="store_true", help="Whether to use tfrecords as dataset")
 
-parser.add_argument("--mxp", default=False, action="store_true",
-                    help="Enable mixed precision")
+parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
 parser.add_argument("--sentence_piece", default=False, action="store_true", help="Whether to use `SentencePiece` model")
 
-parser.add_argument("--device", type=int, default=0,
-                    help="Device's id to run test on")
+parser.add_argument("--device", type=int, default=0, help="Device's id to run test on")
 
-parser.add_argument("--cpu", default=False, action="store_true",
-                    help="Whether to only use cpu")
+parser.add_argument("--cpu", default=False, action="store_true", help="Whether to only use cpu")
 
-parser.add_argument("--subwords", type=str, default=None,
-                    help="Path to file that stores generated subwords")
+parser.add_argument("--subwords", type=str, default=None, help="Path to file that stores generated subwords")
 
-parser.add_argument("--output_name", type=str, default="test",
-                    help="Result filename name prefix")
+parser.add_argument("--output_name", type=str, default="test", help="Result filename name prefix")
 
 args = parser.parse_args()
 
@@ -58,7 +50,7 @@
 setup_devices([args.device], cpu=args.cpu)
 
 from tensorflow_asr.configs.config import Config
-from tensorflow_asr.datasets.asr_dataset import ASRTFRecordTestDataset, ASRSliceTestDataset
+from tensorflow_asr.datasets.asr_dataset import ASRTFRecordDataset, ASRSliceDataset
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import SubwordFeaturizer, SentencePieceFeaturizer
 from tensorflow_asr.runners.base_runners import BaseTester
@@ -80,19 +72,14 @@
 assert args.saved
 
 if args.tfrecords:
-    test_dataset = ASRTFRecordTestDataset(
-        data_paths=config.learning_config.dataset_config.test_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="test", shuffle=False
+    test_dataset = ASRTFRecordDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.test_dataset_config)
     )
 else:
-    test_dataset = ASRSliceTestDataset(
-        data_paths=config.learning_config.dataset_config.test_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="test", shuffle=False
+    test_dataset = ASRSliceDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.test_dataset_config)
     )
 
 # build model

examples/conformer/train_conformer.py

@@ -32,8 +32,6 @@
 
 parser.add_argument("--tfrecords", default=False, action="store_true", help="Whether to use tfrecords")
 
-parser.add_argument("--tfrecords_shards", type=int, default=16, help="Number of tfrecords shards")
-
 parser.add_argument("--tbs", type=int, default=None, help="Train batch size per replica")
 
 parser.add_argument("--ebs", type=int, default=None, help="Evaluation batch size per replica")
@@ -42,10 +40,6 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
-parser.add_argument("--cache", default=False, action="store_true", help="Enable caching for dataset")
-
-parser.add_argument("--bfs", type=int, default=100, help="Buffer size for shuffling")
-
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -66,39 +60,21 @@
 
 if args.tfrecords:
     train_dataset = ASRTFRecordDataset(
-        data_paths=config.learning_config.dataset_config.train_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        augmentations=config.learning_config.augmentations,
-        tfrecords_shards=args.tfrecords_shards,
-        stage="train", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.train_dataset_config)
     )
     eval_dataset = ASRTFRecordDataset(
-        data_paths=config.learning_config.dataset_config.eval_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        tfrecords_shards=args.tfrecords_shards,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="eval", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.eval_dataset_config)
     )
 else:
     train_dataset = ASRSliceDataset(
-        data_paths=config.learning_config.dataset_config.train_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        augmentations=config.learning_config.augmentations,
-        stage="train", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.train_dataset_config)
     )
     eval_dataset = ASRSliceDataset(
-        data_paths=config.learning_config.dataset_config.eval_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="eval", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.eval_dataset_config)
     )
 
 conformer_trainer = TransducerTrainer(

examples/conformer/train_ga_conformer.py

@@ -32,8 +32,6 @@
 
 parser.add_argument("--tfrecords", default=False, action="store_true", help="Whether to use tfrecords")
 
-parser.add_argument("--tfrecords_shards", type=int, default=16, help="Number of tfrecords shards")
-
 parser.add_argument("--tbs", type=int, default=None, help="Train batch size per replica")
 
 parser.add_argument("--ebs", type=int, default=None, help="Evaluation batch size per replica")
@@ -44,10 +42,6 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
-parser.add_argument("--cache", default=False, action="store_true", help="Enable caching for dataset")
-
-parser.add_argument("--bfs", type=int, default=100, help="Buffer size for shuffling")
-
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -68,39 +62,21 @@
 
 if args.tfrecords:
     train_dataset = ASRTFRecordDataset(
-        data_paths=config.learning_config.dataset_config.train_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        augmentations=config.learning_config.augmentations,
-        tfrecords_shards=args.tfrecords_shards,
-        stage="train", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.train_dataset_config)
     )
     eval_dataset = ASRTFRecordDataset(
-        data_paths=config.learning_config.dataset_config.eval_paths,
-        tfrecords_dir=config.learning_config.dataset_config.tfrecords_dir,
-        tfrecords_shards=args.tfrecords_shards,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="eval", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.eval_dataset_config)
     )
 else:
     train_dataset = ASRSliceDataset(
-        data_paths=config.learning_config.dataset_config.train_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        augmentations=config.learning_config.augmentations,
-        stage="train", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.train_dataset_config)
     )
     eval_dataset = ASRSliceDataset(
-        data_paths=config.learning_config.dataset_config.eval_paths,
-        speech_featurizer=speech_featurizer,
-        text_featurizer=text_featurizer,
-        stage="eval", cache=args.cache,
-        shuffle=True, buffer_size=args.bfs,
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.eval_dataset_config)
     )
 
 conformer_trainer = TransducerTrainerGA(