Support for Streaming Conformer Transducer

examples/streaming_conformer/train.py

@@ -0,0 +1,156 @@
+# Copyright 2020 Huy Le Nguyen (@usimarit)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import math
+import argparse
+from tensorflow_asr.utils import env_util
+
+env_util.setup_environment()
+import tensorflow as tf
+physical_devices = tf.config.list_physical_devices('GPU')
+tf.config.experimental.set_memory_growth(physical_devices[0], True)
+
+DEFAULT_YAML = os.path.join(os.path.abspath(os.path.dirname(__file__)), "config.yml")
+
+tf.keras.backend.clear_session()
+
+parser = argparse.ArgumentParser(prog="Conformer Training")
+
+parser.add_argument("--config", type=str, default=DEFAULT_YAML, help="The file path of model configuration file")
+
+parser.add_argument("--max_ckpts", type=int, default=10, help="Max number of checkpoints to keep")
+
+parser.add_argument("--tfrecords", default=False, action="store_true", help="Whether to use tfrecords")
+
+parser.add_argument("--sentence_piece", default=False, action="store_true", help="Whether to use `SentencePiece` model")
+
+parser.add_argument("--subwords", default=False, action="store_true", help="Use subwords")
+
+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")
+
+parser.add_argument("--spx", type=int, default=1, help="Steps per execution for maximizing performance")
+
+parser.add_argument("--metadata", type=str, default=None, help="Path to file containing metadata")
+
+parser.add_argument("--static_length", default=False, action="store_true", help="Use static lengths")
+
+parser.add_argument("--devices", type=int, nargs="*", default=[0], help="Devices' ids to apply distributed training")
+
+parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
+
+args = parser.parse_args()
+
+tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
+
+strategy = env_util.setup_strategy(args.devices)
+
+from tensorflow_asr.configs.config import Config
+from tensorflow_asr.datasets.asr_dataset import ASRMaskedSliceDataset
+from tensorflow_asr.featurizers import speech_featurizers, text_featurizers
+from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer
+from tensorflow_asr.models.transducer.streaming_conformer import StreamingConformer
+from tensorflow_asr.optimizers.schedules import TransformerSchedule
+
+config = Config(args.config)
+speech_featurizer = speech_featurizers.TFSpeechFeaturizer(config.speech_config)
+
+if args.sentence_piece:
+    print("Loading SentencePiece model ...")
+    text_featurizer = text_featurizers.SentencePieceFeaturizer(config.decoder_config)
+elif args.subwords:
+    print("Loading subwords ...")
+    text_featurizer = text_featurizers.SubwordFeaturizer(config.decoder_config)
+else:
+    print("Use characters ...")
+    text_featurizer = text_featurizers.CharFeaturizer(config.decoder_config)
+
+time_reduction_factor = config.model_config['encoder_subsampling']['strides'] * 2
+if args.tfrecords:
+    train_dataset = ASRTFRecordDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.train_dataset_config)
+    )
+    eval_dataset = ASRTFRecordDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        **vars(config.learning_config.eval_dataset_config)
+    )
+else:
+    train_dataset = ASRMaskedSliceDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        time_reduction_factor=time_reduction_factor,
+        **vars(config.learning_config.train_dataset_config)
+    )
+    eval_dataset = ASRMaskedSliceDataset(
+        speech_featurizer=speech_featurizer, text_featurizer=text_featurizer,
+        time_reduction_factor=time_reduction_factor,
+        **vars(config.learning_config.eval_dataset_config)
+    )
+
+train_dataset.load_metadata(args.metadata)
+eval_dataset.load_metadata(args.metadata)
+
+if not args.static_length:
+    speech_featurizer.reset_length()
+    text_featurizer.reset_length()
+
+global_batch_size = args.tbs or config.learning_config.running_config.batch_size
+global_batch_size *= strategy.num_replicas_in_sync
+
+global_eval_batch_size = args.ebs  or global_batch_size
+global_eval_batch_size *= strategy.num_replicas_in_sync
+
+train_data_loader = train_dataset.create(global_batch_size)
+eval_data_loader = eval_dataset.create(global_eval_batch_size)
+
+with strategy.scope():
+    # build model
+    streaming_conformer = StreamingConformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
+    streaming_conformer._build(speech_featurizer.shape)
+    streaming_conformer.summary(line_length=150)
+
+    optimizer = tf.keras.optimizers.Adam(
+        TransformerSchedule(
+            d_model=streaming_conformer.dmodel,
+            warmup_steps=config.learning_config.optimizer_config.pop("warmup_steps", 10000),
+            max_lr=(0.05 / math.sqrt(streaming_conformer.dmodel))
+        ),
+        **config.learning_config.optimizer_config
+    )
+
+    streaming_conformer.compile(
+        optimizer=optimizer,
+        experimental_steps_per_execution=args.spx,
+        global_batch_size=global_batch_size,
+        blank=text_featurizer.blank
+    )
+
+callbacks = [
+    tf.keras.callbacks.ModelCheckpoint(**config.learning_config.running_config.checkpoint),
+    tf.keras.callbacks.experimental.BackupAndRestore(config.learning_config.running_config.states_dir),
+    tf.keras.callbacks.TensorBoard(**config.learning_config.running_config.tensorboard)
+]
+
+streaming_conformer.fit(
+    train_data_loader,
+    batch_size=global_batch_size,
+    epochs=config.learning_config.running_config.num_epochs,
+    steps_per_epoch=train_dataset.total_steps,
+    validation_data=eval_data_loader,
+    validation_batch_size=global_eval_batch_size,
+    validation_steps=eval_dataset.total_steps,
+    callbacks=callbacks,
+)

tensorflow_asr/datasets/asr_dataset.py

@@ -366,3 +366,165 @@ def create(self, batch_size: int):
         dataset = tf.data.Dataset.from_tensor_slices(self.entries)
         dataset = dataset.map(self.load, num_parallel_calls=AUTOTUNE)
         return self.process(dataset, batch_size)
+
+
+class ASRMaskedSliceDataset(ASRSliceDataset):
+    """ Dataset for ASR with rolling mask """
+
+    def __init__(self,
+                 stage: str,
+                 speech_featurizer: SpeechFeaturizer,
+                 text_featurizer: TextFeaturizer,
+                 data_paths: list,
+                 augmentations: Augmentation = Augmentation(None),
+                 cache: bool = False,
+                 shuffle: bool = False,
+                 indefinite: bool = False,
+                 drop_remainder: bool = True,
+                 use_tf: bool = False,
+                 buffer_size: int = BUFFER_SIZE,
+                 history_window_size: int = 3,
+                 input_chunk_duration: int = 250,
+                 time_reduction_factor: int = 4,
+                 **kwargs):
+        super(ASRMaskedSliceDataset, self).__init__(
+            data_paths=data_paths, augmentations=augmentations,
+            cache=cache, shuffle=shuffle, stage=stage, buffer_size=buffer_size,
+            drop_remainder=drop_remainder, use_tf=use_tf, indefinite=indefinite,
+            speech_featurizer=speech_featurizer, text_featurizer=text_featurizer
+        )
+        self.speech_featurizer = speech_featurizer
+        self.text_featurizer = text_featurizer
+        self.history_window_size = history_window_size
+        self.input_chunk_size = input_chunk_duration * self.speech_featurizer.sample_rate // 1000
+        self.time_reduction_factor = time_reduction_factor
+
+    def calculate_mask(self, num_frames):
+        frame_step = self.speech_featurizer.frame_step
+        frames_per_chunk = self.input_chunk_size // frame_step
+
+        num_frames = tf.cast(tf.math.ceil(num_frames / self.time_reduction_factor), tf.int32)
+
+        def _calculate_mask(num_frames, frames_per_chunk, history_window_size):
+            mask = np.zeros((num_frames, num_frames), dtype=np.int32)
+            for i in range(num_frames):
+                # Frames in the same chunk can see each other
+                # If frames in `history_window_size` are in other chunks, the full chunks are visible
+                current_chunk_index = i // frames_per_chunk
+                history_chunk_index = (i - history_window_size) // frames_per_chunk
+                for curr in range(history_chunk_index, current_chunk_index + 1):
+                    for j in range(frames_per_chunk):
+                        base_index = curr * frames_per_chunk
+                        if base_index + j < 0 or base_index + j >= num_frames:
+                            continue
+                        mask[i, base_index + j] = 1
+            return mask
+
+        return tf.numpy_function(
+            _calculate_mask, inp=[num_frames, frames_per_chunk, self.history_window_size], Tout=tf.int32
+        )
+
+    def preprocess(self, path: tf.Tensor, audio: tf.Tensor, indices: tf.Tensor):
+        preprocessed_inputs = super(ASRMaskedSliceDataset, self).preprocess(path, audio, indices)
+
+        input_length = preprocessed_inputs[2]
+        mask = self.calculate_mask(input_length)
+        mask.set_shape((None, None))
+
+        return (*preprocessed_inputs, mask)
+
+    def tf_preprocess(self, path: tf.Tensor, audio: tf.Tensor, indices: tf.Tensor):
+        preprocessed_inputs = super(ASRMaskedSliceDataset, self).tf_preprocess(path, audio, indices)
+
+        input_length = preprocessed_inputs[2]
+        mask = self.calculate_mask(input_length)
+        mask.set_shape((None, None))
+
+        return (*preprocessed_inputs, mask)
+
+    # -------------------------------- CREATION -------------------------------------
+
+    def parse(self, path: tf.Tensor, audio: tf.Tensor, indices: tf.Tensor):
+        """
+        Returns:
+            path, features, input_lengths, labels, label_lengths, pred_inp, mask
+        """
+        if self.use_tf: data = self.tf_preprocess(path, audio, indices)
+        else: data = self.preprocess(path, audio, indices)
+
+        _, features, input_length, label, label_length, prediction, prediction_length, mask = data
+
+        return (
+            data_util.create_inputs(
+                inputs=features,
+                inputs_length=input_length,
+                predictions=prediction,
+                predictions_length=prediction_length,
+                mask = mask
+            ),
+            data_util.create_labels(
+                labels=label,
+                labels_length=label_length
+            )
+        )
+
+    def process(self, dataset: tf.data.Dataset, batch_size: int):
+        dataset = dataset.map(self.parse, num_parallel_calls=AUTOTUNE)
+        self.total_steps = math_util.get_num_batches(self.total_steps, batch_size, drop_remainders=self.drop_remainder)
+
+        if self.cache:
+            dataset = dataset.cache()
+
+        if self.shuffle:
+            dataset = dataset.shuffle(self.buffer_size, reshuffle_each_iteration=True)
+
+        if self.indefinite and self.total_steps:
+            dataset = dataset.repeat()
+
+        # PADDED BATCH the dataset
+        dataset = dataset.padded_batch(
+            batch_size=batch_size,
+            padded_shapes=(
+                data_util.create_inputs(
+                    inputs=tf.TensorShape(self.speech_featurizer.shape),
+                    inputs_length=tf.TensorShape([]),
+                    predictions=tf.TensorShape(self.text_featurizer.prepand_shape),
+                    predictions_length=tf.TensorShape([]),
+                    mask=tf.TensorShape([self.speech_featurizer.shape[0], self.speech_featurizer.shape[0]])
+                ),
+                data_util.create_labels(
+                    labels=tf.TensorShape(self.text_featurizer.shape),
+                    labels_length=tf.TensorShape([])
+                ),
+            ),
+            padding_values=(
+                data_util.create_inputs(
+                    inputs= 0.,
+                    inputs_length=0,
+                    predictions=self.text_featurizer.blank,
+                    predictions_length=0,
+                    mask=0
+                ),
+                data_util.create_labels(
+                    labels=self.text_featurizer.blank,
+                    labels_length=0
+                )
+            ),
+            drop_remainder = self.drop_remainder
+        )
+
+        # PREFETCH to improve speed of input length
+        dataset = dataset.prefetch(AUTOTUNE)
+        return dataset
+
+    def create(self, batch_size: int):
+        self.read_entries()
+        if not self.total_steps or self.total_steps == 0: return None
+        dataset = tf.data.Dataset.from_generator(
+            self.generator,
+            output_types=(tf.string, tf.string, tf.string),
+            output_shapes=(tf.TensorShape([]), tf.TensorShape([]), tf.TensorShape([]))
+        )
+        return self.process(dataset, batch_size)
+
+# TODO: Create masked TFRecords dataset