✍️ update testing functions and scripts

Author: nglehuy

examples/conformer/test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from tqdm import tqdm
 import argparse
 from tensorflow_asr.utils import env_util, file_util
 
@@ -58,6 +59,7 @@
 from tensorflow_asr.featurizers.speech_featurizers import TFSpeechFeaturizer
 from tensorflow_asr.featurizers.text_featurizers import SubwordFeaturizer, SentencePieceFeaturizer, CharFeaturizer
 from tensorflow_asr.models.transducer.conformer import Conformer
+from tensorflow_asr.utils import app_util
 
 config = Config(args.config)
 speech_featurizer = TFSpeechFeaturizer(config.speech_config)
@@ -97,13 +99,20 @@
 batch_size = args.bs or config.learning_config.running_config.batch_size
 test_data_loader = test_dataset.create(batch_size)
 
-results = conformer.predict(test_data_loader)
-
 with file_util.save_file(file_util.preprocess_paths(args.output)) as filepath:
-    print(f"Saving result to {args.output} ...")
-    with open(filepath, "w") as openfile:
-        openfile.write("PATH\tDURATION\tGROUNDTRUTH\tGREEDY\tBEAMSEARCH\n")
-        for i, entry in test_dataset.entries:
-            groundtruth, greedy, beamsearch = results[i]
-            path, duration, _ = entry
-            openfile.write(f"{path}\t{duration}\t{groundtruth}\t{greedy}\t{beamsearch}\n")
+    overwrite = False
+    if tf.io.gfile.exists(filepath):
+        overwrite = input("Overwrite existing result file? (y/n): ").lower() == "y"
+    if overwrite:
+        results = conformer.predict(test_data_loader, verbose=1)
+        print(f"Saving result to {args.output} ...")
+        with open(filepath, "w") as openfile:
+            openfile.write("PATH\tDURATION\tGROUNDTRUTH\tGREEDY\tBEAMSEARCH\n")
+            progbar = tqdm(total=test_dataset.total_steps, unit="batch")
+            for i, pred in enumerate(results):
+                groundtruth, greedy, beamsearch = [x.decode('utf-8') for x in pred]
+                path, duration, _ = test_dataset.entries[i]
+                openfile.write(f"{path}\t{duration}\t{groundtruth}\t{greedy}\t{beamsearch}\n")
+                progbar.update(1)
+            progbar.close()
+    app_util.evaluate_results(filepath)

tensorflow_asr/metrics/error_rates.py

@@ -30,4 +30,4 @@ def update_state(self, decode: tf.Tensor, target: tf.Tensor):
         self.denominator.assign_add(d)
 
     def result(self):
-        return tf.math.divide_no_nan(self.numerator, self.denominator) * 100
+        return tf.math.divide_no_nan(self.numerator, self.denominator)

tensorflow_asr/models/base_model.py

@@ -19,6 +19,10 @@
 
 
 class BaseModel(tf.keras.Model):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._metrics = {}
+
     def save(self,
              filepath,
              overwrite=True,
@@ -66,7 +70,10 @@ def load_weights(self,
 
     @property
     def metrics(self):
-        return [self.loss_metric]
+        return self._metrics.values()
+
+    def add_metric(self, metric: tf.keras.metrics.Metric):
+        self._metrics.append({metric.name: metric})
 
     def _build(self, *args, **kwargs):
         raise NotImplementedError()
@@ -76,7 +83,8 @@ def compile(self, loss, optimizer, run_eagerly=None, **kwargs):
         if not env_util.has_tpu():
             optimizer = mxp.experimental.LossScaleOptimizer(tf.keras.optimizers.get(optimizer), "dynamic")
             self.use_loss_scale = True
-        self.loss_metric = tf.keras.metrics.Mean(name="loss", dtype=tf.float32)
+        loss_metric = tf.keras.metrics.Mean(name="loss", dtype=tf.float32)
+        self._metrics = {loss_metric.name: loss_metric}
         super().compile(optimizer=optimizer, loss=loss, run_eagerly=run_eagerly, **kwargs)
 
     # -------------------------------- STEP FUNCTIONS -------------------------------------
@@ -92,14 +100,14 @@ def train_step(self, batch):
         if self.use_loss_scale:
             gradients = self.optimizer.get_unscaled_gradients(gradients)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        self.loss_metric.update_state(loss)
+        self._metrics["loss"].update_state(loss)
         return {m.name: m.result() for m in self.metrics}
 
     def test_step(self, batch):
         inputs, y_true = batch
         y_pred = self(inputs, training=False)
         loss = self.loss(y_true, y_pred)
-        self.loss_metric.update_state(loss)
+        self._metrics["loss"].update_state(loss)
         return {m.name: m.result() for m in self.metrics}
 
     def predict_step(self, batch):

tensorflow_asr/utils/app_util.py

@@ -0,0 +1,45 @@
+# 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.
+
+from tqdm import tqdm
+import tensorflow as tf
+
+from .metric_util import wer, cer
+from ..metrics.error_rates import ErrorRate
+from .file_util import read_file
+
+
+def evaluate_results(filepath: str):
+    print(f"Evaluating result from {filepath} ...")
+    metrics = {
+        "greedy_wer": ErrorRate(wer, name="greedy_wer", dtype=tf.float32),
+        "greedy_cer": ErrorRate(cer, name="greedy_cer", dtype=tf.float32),
+        "beamsearch_wer": ErrorRate(wer, name="beamsearch_wer", dtype=tf.float32),
+        "beamsearch_cer": ErrorRate(cer, name="beamsearch_cer", dtype=tf.float32)
+    }
+    with read_file(filepath) as path:
+        with open(path, "r", encoding="utf-8") as openfile:
+            lines = openfile.read().splitlines()
+            lines = lines[1:]  # skip header
+    for eachline in tqdm(lines):
+        _, _, groundtruth, greedy, beamsearch = eachline.split("\t")
+        groundtruth = tf.convert_to_tensor([groundtruth], dtype=tf.string)
+        greedy = tf.convert_to_tensor([greedy], dtype=tf.string)
+        beamsearch = tf.convert_to_tensor([beamsearch], dtype=tf.string)
+        metrics["greedy_wer"].update_state(decode=greedy, target=groundtruth)
+        metrics["greedy_cer"].update_state(decode=greedy, target=groundtruth)
+        metrics["beamsearch_wer"].update_state(decode=beamsearch, target=groundtruth)
+        metrics["beamsearch_cer"].update_state(decode=beamsearch, target=groundtruth)
+    for key, value in metrics.items():
+        print(f"{key}: {value.result().numpy()}")