✍️ update metrics for training using builtin keras

Author: nglehuy

tensorflow_asr/losses/keras/ctc_losses.py

@@ -13,13 +13,11 @@
 # limitations under the License.
 
 import tensorflow as tf
-from tensorflow.python.keras.utils import losses_utils
-
 from .. import ctc_loss
 
 
 class CtcLoss(tf.keras.losses.Loss):
-    def __init__(self, blank=0, global_batch_size=None, reduction=losses_utils.ReductionV2.NONE, name=None):
+    def __init__(self, blank=0, global_batch_size=None, reduction=tf.keras.losses.Reduction.NONE, name=None):
         super(CtcLoss, self).__init__(reduction=reduction, name=name)
         self.blank = blank
         self.global_batch_size = global_batch_size

tensorflow_asr/losses/keras/rnnt_losses.py

@@ -13,13 +13,11 @@
 # limitations under the License.
 
 import tensorflow as tf
-from tensorflow.python.keras.utils import losses_utils
-
 from .. import rnnt_loss
 
 
 class RnntLoss(tf.keras.losses.Loss):
-    def __init__(self, blank=0, global_batch_size=None, reduction=losses_utils.ReductionV2.NONE, name=None):
+    def __init__(self, blank=0, global_batch_size=None, reduction=tf.keras.losses.Reduction.NONE, name=None):
         super(RnntLoss, self).__init__(reduction=reduction, name=name)
         self.blank = blank
         self.global_batch_size = global_batch_size

tensorflow_asr/models/keras/ctc.py

@@ -22,27 +22,25 @@
 
 class CtcModel(BaseCtcModel):
     """ Keras CTC Model Warper """
+    @property
+    def metrics(self):
+        return [self.loss_metric]
 
-    def compile(self, optimizer, global_batch_size, blank=0, use_loss_scale=False,
-                loss_weights=None, weighted_metrics=None, run_eagerly=None, **kwargs):
+    def compile(self, optimizer, global_batch_size, blank=0, use_loss_scale=False, run_eagerly=None, **kwargs):
         loss = CtcLoss(blank=blank, global_batch_size=global_batch_size)
         self.use_loss_scale = use_loss_scale
         if self.use_loss_scale:
-            optimizer = mxp.experimental.LossScaleOptimizer(tf.keras.optimizers.get(optimizer), 'dynamic')
-        super(CtcModel, self).compile(
-            optimizer=optimizer, loss=loss,
-            loss_weights=loss_weights, weighted_metrics=weighted_metrics,
-            run_eagerly=run_eagerly,
-            **kwargs
-        )
+            optimizer = mxp.experimental.LossScaleOptimizer(tf.keras.optimizers.get(optimizer), "dynamic")
+        self.loss_metric = tf.keras.metrics.Mean(name="ctc_loss", dtype=tf.float32)
+        super(CtcModel, self).compile(optimizer=optimizer, loss=loss, run_eagerly=run_eagerly, **kwargs)
 
     def train_step(self, batch):
         x, y_true = batch
         with tf.GradientTape() as tape:
-            logit = self(x['input'], training=True)
+            logit = self(x["input"], training=True)
             y_pred = {
-                'logit': logit,
-                'logit_length': get_reduced_length(x['input_length'], self.time_reduction_factor)
+                "logit": logit,
+                "logit_length": get_reduced_length(x["input_length"], self.time_reduction_factor)
             }
             loss = self.loss(y_true, y_pred)
             if self.use_loss_scale:
@@ -53,14 +51,16 @@ def train_step(self, batch):
         else:
             gradients = tape.gradient(loss, self.trainable_weights)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        return {"ctc_loss": loss}
+        self.loss_metric.update_state(loss)
+        return {m.name: m.result() for m in self.metrics}
 
     def test_step(self, batch):
         x, y_true = batch
         logit = self(x, training=False)
         y_pred = {
-            'logit': logit,
-            'logit_length': get_reduced_length(x['input_length'], self.time_reduction_factor)
+            "logit": logit,
+            "logit_length": get_reduced_length(x["input_length"], self.time_reduction_factor)
         }
         loss = self.loss(y_true, y_pred)
-        return {"ctc_loss": loss}
+        self.loss_metric.update_state(loss)
+        return {m.name: m.result() for m in self.metrics}

tensorflow_asr/models/keras/transducer.py

@@ -23,6 +23,9 @@
 
 class Transducer(BaseTransducer):
     """ Keras Transducer Model Warper """
+    @property
+    def metrics(self):
+        return [self.loss_metric]
 
     def _build(self, input_shape, prediction_shape=[None], batch_size=None):
         inputs = tf.keras.Input(shape=input_shape, batch_size=batch_size, dtype=tf.float32)
@@ -48,18 +51,13 @@ def call(self, inputs, training=False, **kwargs):
             "logit_length": get_reduced_length(inputs["input_length"], self.time_reduction_factor)
         }
 
-    def compile(self, optimizer, global_batch_size, blank=0, use_loss_scale=False,
-                loss_weights=None, weighted_metrics=None, run_eagerly=None, **kwargs):
+    def compile(self, optimizer, global_batch_size, blank=0, use_loss_scale=False, run_eagerly=None, **kwargs):
         loss = RnntLoss(blank=blank, global_batch_size=global_batch_size)
         self.use_loss_scale = use_loss_scale
         if self.use_loss_scale:
-            optimizer = mxp.experimental.LossScaleOptimizer(tf.keras.optimizers.get(optimizer), 'dynamic')
-        super(Transducer, self).compile(
-            optimizer=optimizer, loss=loss,
-            loss_weights=loss_weights, weighted_metrics=weighted_metrics,
-            run_eagerly=run_eagerly,
-            **kwargs
-        )
+            optimizer = mxp.experimental.LossScaleOptimizer(tf.keras.optimizers.get(optimizer), "dynamic")
+        self.loss_metric = tf.keras.metrics.Mean(name="rnnt_loss", dtype=tf.float32)
+        super(Transducer, self).compile(optimizer=optimizer, loss=loss, run_eagerly=run_eagerly, **kwargs)
 
     def train_step(self, batch):
         x, y_true = batch
@@ -79,7 +77,8 @@ def train_step(self, batch):
         else:
             gradients = tape.gradient(loss, self.trainable_weights)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        return {"rnnt_loss": loss}
+        self.loss_metric.update_state(loss)
+        return {m.name: m.result() for m in self.metrics}
 
     def test_step(self, batch):
         x, y_true = batch
@@ -90,4 +89,5 @@ def test_step(self, batch):
             "prediction_length": x["prediction_length"],
         }, training=False)
         loss = self.loss(y_true, y_pred)
-        return {"rnnt_loss": loss}
+        self.loss_metric.update_state(loss)
+        return {m.name: m.result() for m in self.metrics}