✍️ update keras return dict

Author: nglehuy

examples/conformer/config.yml

@@ -58,8 +58,9 @@ model_config:
   prediction_rnn_type: lstm
   prediction_rnn_implementation: 2
   prediction_layer_norm: False
-  prediction_projection_units: 0
+  prediction_projection_units: 144
   joint_dim: 640
+  prejoint_linear: False
   joint_activation: tanh
 
 learning_config:
@@ -112,7 +113,7 @@ learning_config:
   running_config:
     batch_size: 2
     accumulation_steps: 4
-    num_epochs: 20
+    num_epochs: 50
     outdir: /mnt/Miscellanea/Models/local/conformer
     log_interval_steps: 300
     eval_interval_steps: 500

tensorflow_asr/models/keras/ctc.py

@@ -47,7 +47,7 @@ def train_step(self, batch):
         scaled_gradients = tape.gradient(scaled_loss, self.trainable_weights)
         gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        return {"train_ctc_loss": loss}
+        return {"ctc_loss": loss}
 
     def test_step(self, batch):
         x, y_true = batch
@@ -57,4 +57,4 @@ def test_step(self, batch):
             'logit_length': get_reduced_length(x['input_length'], self.time_reduction_factor)
         }
         loss = self.loss(y_true, y_pred)
-        return {"val_ctc_loss": loss}
+        return {"ctc_loss": loss}

tensorflow_asr/models/keras/transducer.py

@@ -73,7 +73,7 @@ def train_step(self, batch):
         scaled_gradients = tape.gradient(scaled_loss, self.trainable_weights)
         gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-        return {"train_rnnt_loss": loss}
+        return {"rnnt_loss": loss}
 
     def test_step(self, batch):
         x, y_true = batch
@@ -84,4 +84,4 @@ def test_step(self, batch):
             "prediction_length": x["prediction_length"],
         }, training=False)
         loss = self.loss(y_true, y_pred)
-        return {"val_rnnt_loss": loss}
+        return {"rnnt_loss": loss}