✍️ update examples train scripts

nglehuy · nglehuy · commit 2c90f67f284b · 2021-05-02T13:41:01.000+07:00
diff --git a/examples/conformer/test.py b/examples/conformer/test.py
@@ -83,7 +83,7 @@
 # build model
 conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer.make(speech_featurizer.shape)
-conformer.load_weights(args.saved)
+conformer.load_weights(args.saved, by_name=True)
 conformer.summary(line_length=100)
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/conformer/tflite.py b/examples/conformer/tflite.py
@@ -53,7 +53,7 @@
 # build model
 conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer.make(speech_featurizer.shape)
-conformer.load_weights(args.saved)
+conformer.load_weights(args.saved, by_name=True)
 conformer.summary(line_length=100)
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/conformer/train.py b/examples/conformer/train.py
@@ -46,6 +46,8 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
+parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model")
+
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -119,8 +121,9 @@
         prediction_shape=text_featurizer.prepand_shape,
         batch_size=global_batch_size
     )
+    if args.pretrained:
+        conformer.load_weights(args.pretrained, by_name=True, skip_mismatch=True)
     conformer.summary(line_length=100)
-
     optimizer = tf.keras.optimizers.Adam(
         TransformerSchedule(
             d_model=conformer.dmodel,
@@ -129,7 +132,6 @@
         ),
         **config.learning_config.optimizer_config
     )
-
     conformer.compile(
         optimizer=optimizer,
         experimental_steps_per_execution=args.spx,
diff --git a/examples/contextnet/test.py b/examples/contextnet/test.py
@@ -83,7 +83,7 @@
 # build model
 contextnet = ContextNet(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 contextnet.make(speech_featurizer.shape)
-contextnet.load_weights(args.saved)
+contextnet.load_weights(args.saved, by_name=True)
 contextnet.summary(line_length=100)
 contextnet.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/contextnet/tflite.py b/examples/contextnet/tflite.py
@@ -53,7 +53,7 @@
 # build model
 contextnet = ContextNet(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 contextnet.make(speech_featurizer.shape)
-contextnet.load_weights(args.saved)
+contextnet.load_weights(args.saved, by_name=True)
 contextnet.summary(line_length=100)
 contextnet.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/contextnet/train.py b/examples/contextnet/train.py
@@ -46,6 +46,8 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
+parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model")
+
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -119,8 +121,9 @@
         prediction_shape=text_featurizer.prepand_shape,
         batch_size=global_batch_size
     )
+    if args.pretrained:
+        contextnet.load_weights(args.pretrained, by_name=True, skip_mismatch=True)
     contextnet.summary(line_length=100)
-
     optimizer = tf.keras.optimizers.Adam(
         TransformerSchedule(
             d_model=contextnet.dmodel,
@@ -129,7 +132,6 @@
         ),
         **config.learning_config.optimizer_config
     )
-
     contextnet.compile(
         optimizer=optimizer,
         experimental_steps_per_execution=args.spx,
diff --git a/examples/deepspeech2/test.py b/examples/deepspeech2/test.py
@@ -83,7 +83,7 @@
 # build model
 deepspeech2 = DeepSpeech2(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 deepspeech2.make(speech_featurizer.shape)
-deepspeech2.load_weights(args.saved)
+deepspeech2.load_weights(args.saved, by_name=True)
 deepspeech2.summary(line_length=100)
 deepspeech2.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/deepspeech2/tflite.py b/examples/deepspeech2/tflite.py
@@ -53,7 +53,7 @@
 # build model
 deepspeech2 = DeepSpeech2(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 deepspeech2.make(speech_featurizer.shape)
-deepspeech2.load_weights(args.saved)
+deepspeech2.load_weights(args.saved, by_name=True)
 deepspeech2.summary(line_length=100)
 deepspeech2.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/deepspeech2/train.py b/examples/deepspeech2/train.py
@@ -45,6 +45,8 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
+parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model")
+
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -113,6 +115,8 @@
     # build model
     deepspeech2 = DeepSpeech2(**config.model_config, vocabulary_size=text_featurizer.num_classes)
     deepspeech2.make(speech_featurizer.shape, batch_size=global_batch_size)
+    if args.pretrained:
+        deepspeech2.load_weights(args.pretrained, by_name=True, skip_mismatch=True)
     deepspeech2.summary(line_length=100)
     deepspeech2.compile(
         optimizer=config.learning_config.optimizer_config,
diff --git a/examples/demonstration/conformer.py b/examples/demonstration/conformer.py
@@ -64,7 +64,7 @@
 # build model
 conformer = Conformer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 conformer.make(speech_featurizer.shape)
-conformer.load_weights(args.saved)
+conformer.load_weights(args.saved, by_name=True, skip_mismatch=True)
 conformer.summary(line_length=120)
 conformer.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/demonstration/rnn_transducer.py b/examples/demonstration/rnn_transducer.py
@@ -0,0 +1,99 @@
+# 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 argparse
+from tensorflow_asr.utils import env_util, math_util, data_util
+
+env_util.setup_environment()
+import tensorflow as tf
+
+parser = argparse.ArgumentParser(prog="Rnn Transducer non streaming")
+
+parser.add_argument("filename", metavar="FILENAME", help="audio file to be played back")
+
+parser.add_argument("--config", type=str, default=None, help="Path to rnnt config yaml")
+
+parser.add_argument("--saved", type=str, default=None, help="Path to rnnt saved h5 weights")
+
+parser.add_argument("--beam_width", type=int, default=0, help="Beam width")
+
+parser.add_argument("--timestamp", default=False, action="store_true", help="Return with timestamp")
+
+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("--subwords", default=False, action="store_true", help="Path to file that stores generated subwords")
+
+parser.add_argument("--sentence_piece", default=False, action="store_true", help="Whether to use `SentencePiece` model")
+
+args = parser.parse_args()
+
+env_util.setup_devices([args.device], cpu=args.cpu)
+
+from tensorflow_asr.configs.config import Config
+from tensorflow_asr.featurizers.speech_featurizers import read_raw_audio
+from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
+from tensorflow_asr.featurizers.text_featurizers import CharFeaturizer, SubwordFeaturizer, SentencePieceFeaturizer
+from tensorflow_asr.models.transducer.rnn_transducer import RnnTransducer
+
+config = Config(args.config)
+speech_featurizer = TFSpeechFeaturizer(config.speech_config)
+if args.sentence_piece:
+    print("Loading SentencePiece model ...")
+    text_featurizer = SentencePieceFeaturizer(config.decoder_config)
+elif args.subwords:
+    print("Loading subwords ...")
+    text_featurizer = SubwordFeaturizer(config.decoder_config)
+else:
+    text_featurizer = CharFeaturizer(config.decoder_config)
+text_featurizer.decoder_config.beam_width = args.beam_width
+
+# build model
+rnnt = RnnTransducer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
+rnnt.make(speech_featurizer.shape)
+rnnt.load_weights(args.saved, by_name=True, skip_mismatch=True)
+rnnt.summary(line_length=120)
+rnnt.add_featurizers(speech_featurizer, text_featurizer)
+
+signal = read_raw_audio(args.filename)
+features = speech_featurizer.tf_extract(signal)
+input_length = math_util.get_reduced_length(tf.shape(features)[0], rnnt.time_reduction_factor)
+
+if args.beam_width:
+    transcript = rnnt.recognize_beam(
+        data_util.create_inputs(
+            inputs=features[None, ...],
+            inputs_length=input_length[None, ...]
+        )
+    )
+    print("Transcript:", transcript[0].numpy().decode("UTF-8"))
+elif args.timestamp:
+    transcript, stime, etime, _, _, _ = rnnt.recognize_tflite_with_timestamp(
+        signal=signal,
+        predicted=tf.constant(text_featurizer.blank, dtype=tf.int32),
+        encoder_states=rnnt.encoder.get_initial_state(),
+        prediction_states=rnnt.predict_net.get_initial_state()
+    )
+    print("Transcript:", transcript)
+    print("Start time:", stime)
+    print("End time:", etime)
+else:
+    transcript = rnnt.recognize(
+        data_util.create_inputs(
+            inputs=features[None, ...],
+            inputs_length=input_length[None, ...]
+        )
+    )
+    print("Transcript:", transcript[0].numpy().decode("UTF-8"))
diff --git a/examples/jasper/test.py b/examples/jasper/test.py
@@ -83,7 +83,7 @@
 # build model
 jasper = Jasper(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 jasper.make(speech_featurizer.shape)
-jasper.load_weights(args.saved)
+jasper.load_weights(args.saved, by_name=True)
 jasper.summary(line_length=100)
 jasper.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/jasper/tflite.py b/examples/jasper/tflite.py
@@ -53,7 +53,7 @@
 # build model
 jasper = Jasper(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 jasper.make(speech_featurizer.shape)
-jasper.load_weights(args.saved)
+jasper.load_weights(args.saved, by_name=True)
 jasper.summary(line_length=100)
 jasper.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/jasper/train.py b/examples/jasper/train.py
@@ -45,6 +45,8 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
+parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model")
+
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -113,6 +115,8 @@
     # build model
     jasper = Jasper(**config.model_config, vocabulary_size=text_featurizer.num_classes)
     jasper.make(speech_featurizer.shape, batch_size=global_batch_size)
+    if args.pretrained:
+        jasper.load_weights(args.pretrained, by_name=True, skip_mismatch=True)
     jasper.summary(line_length=100)
     jasper.compile(
         optimizer=config.learning_config.optimizer_config,
diff --git a/examples/rnn_transducer/test.py b/examples/rnn_transducer/test.py
@@ -83,7 +83,7 @@
 # build model
 rnn_transducer = RnnTransducer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 rnn_transducer.make(speech_featurizer.shape)
-rnn_transducer.load_weights(args.saved)
+rnn_transducer.load_weights(args.saved, by_name=True)
 rnn_transducer.summary(line_length=100)
 rnn_transducer.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/rnn_transducer/tflite.py b/examples/rnn_transducer/tflite.py
@@ -53,7 +53,7 @@
 # build model
 rnn_transducer = RnnTransducer(**config.model_config, vocabulary_size=text_featurizer.num_classes)
 rnn_transducer.make(speech_featurizer.shape)
-rnn_transducer.load_weights(args.saved)
+rnn_transducer.load_weights(args.saved, by_name=True)
 rnn_transducer.summary(line_length=100)
 rnn_transducer.add_featurizers(speech_featurizer, text_featurizer)
 
diff --git a/examples/rnn_transducer/train.py b/examples/rnn_transducer/train.py
@@ -46,6 +46,8 @@
 
 parser.add_argument("--mxp", default=False, action="store_true", help="Enable mixed precision")
 
+parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model")
+
 args = parser.parse_args()
 
 tf.config.optimizer.set_experimental_options({"auto_mixed_precision": args.mxp})
@@ -119,6 +121,8 @@
         prediction_shape=text_featurizer.prepand_shape,
         batch_size=global_batch_size
     )
+    if args.pretrained:
+        rnn_transducer.load_weights(args.pretrained, by_name=True, skip_mismatch=True)
     rnn_transducer.summary(line_length=100)
     optimizer = tf.keras.optimizers.Adam(
         TransformerSchedule(
diff --git a/tensorflow_asr/featurizers/speech_featurizers.py b/tensorflow_asr/featurizers/speech_featurizers.py
@@ -16,6 +16,7 @@
 import io
 import abc
 import math
+from typing import Union
 import numpy as np
 import librosa
 import soundfile as sf
@@ -39,7 +40,7 @@ def load_and_convert_to_wav(path: str) -> tf.Tensor:
     return tf.audio.encode_wav(tf.expand_dims(wave, axis=-1), sample_rate=rate)
 
 
-def read_raw_audio(audio, sample_rate=16000):
+def read_raw_audio(audio: Union[str, bytes, np.ndarray], sample_rate=16000) -> np.ndarray:
     if isinstance(audio, str):
         wave, _ = librosa.load(os.path.expanduser(audio), sr=sample_rate, mono=True)
     elif isinstance(audio, bytes):
@@ -55,7 +56,7 @@ def read_raw_audio(audio, sample_rate=16000):
     return wave
 
 
-def tf_read_raw_audio(audio: tf.Tensor, sample_rate=16000):
+def tf_read_raw_audio(audio: tf.Tensor, sample_rate=16000) -> tf.Tensor:
     wave, rate = tf.audio.decode_wav(audio, desired_channels=1, desired_samples=-1)
     if not env_util.has_devices("TPU"):
         resampled = tfio.audio.resample(wave, rate_in=tf.cast(rate, dtype=tf.int64), rate_out=sample_rate)
@@ -91,7 +92,7 @@ def merge_slices(slices: np.ndarray) -> np.ndarray:
     return np.reshape(slices, [-1])
 
 
-def normalize_audio_feature(audio_feature: np.ndarray, per_frame=False):
+def normalize_audio_feature(audio_feature: np.ndarray, per_frame=False) -> np.ndarray:
     """ Mean and variance normalization """
     axis = 1 if per_frame else None
     mean = np.mean(audio_feature, axis=axis)
@@ -100,7 +101,7 @@ def normalize_audio_feature(audio_feature: np.ndarray, per_frame=False):
     return normalized
 
 
-def tf_normalize_audio_features(audio_feature: tf.Tensor, per_frame=False):
+def tf_normalize_audio_features(audio_feature: tf.Tensor, per_frame=False) -> tf.Tensor:
     """
     TF Mean and variance features normalization
     Args:
@@ -115,13 +116,13 @@ def tf_normalize_audio_features(audio_feature: tf.Tensor, per_frame=False):
     return (audio_feature - mean) / std_dev
 
 
-def normalize_signal(signal: np.ndarray):
+def normalize_signal(signal: np.ndarray) -> np.ndarray:
     """ Normailize signal to [-1, 1] range """
     gain = 1.0 / (np.max(np.abs(signal)) + 1e-9)
     return signal * gain
 
 
-def tf_normalize_signal(signal: tf.Tensor):
+def tf_normalize_signal(signal: tf.Tensor) -> tf.Tensor:
     """
     TF Normailize signal to [-1, 1] range
     Args:
@@ -134,7 +135,7 @@ def tf_normalize_signal(signal: tf.Tensor):
     return signal * gain
 
 
-def preemphasis(signal: np.ndarray, coeff=0.97):
+def preemphasis(signal: np.ndarray, coeff=0.97) -> np.ndarray:
     if not coeff or coeff <= 0.0:
         return signal
     return np.append(signal[0], signal[1:] - coeff * signal[:-1])
@@ -156,7 +157,7 @@ def tf_preemphasis(signal: tf.Tensor, coeff=0.97):
     return tf.concat([s0, s1], axis=-1)
 
 
-def depreemphasis(signal: np.ndarray, coeff=0.97):
+def depreemphasis(signal: np.ndarray, coeff=0.97) -> np.ndarray:
     if not coeff or coeff <= 0.0: return signal
     x = np.zeros(signal.shape[0], dtype=np.float32)
     x[0] = signal[0]
@@ -165,7 +166,7 @@ def depreemphasis(signal: np.ndarray, coeff=0.97):
     return x
 
 
-def tf_depreemphasis(signal: tf.Tensor, coeff=0.97):
+def tf_depreemphasis(signal: tf.Tensor, coeff=0.97) -> tf.Tensor:
     """
     TF Depreemphasis
     Args:
diff --git a/tensorflow_asr/models/transducer/rnn_transducer.py b/tensorflow_asr/models/transducer/rnn_transducer.py
diff --git a/tensorflow_asr/models/transducer/transducer.py b/tensorflow_asr/models/transducer/transducer.py
diff --git a/tests/featurizer/test_speech_featurizer.ipynb b/tests/featurizer/test_speech_featurizer.ipynb
