🚀 add post joint linear

Author: nglehuy

tensorflow_asr/models/conformer.py

@@ -391,6 +391,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer=L2,
@@ -428,6 +429,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/contextnet.py

@@ -210,6 +210,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer=L2,
@@ -238,6 +239,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/keras/conformer.py

@@ -43,6 +43,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer=L2,
@@ -80,6 +81,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/keras/contextnet.py

@@ -38,6 +38,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer=L2,
@@ -66,6 +67,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/keras/streaming_transducer.py

@@ -40,6 +40,7 @@ def __init__(self,
                  joint_dim: int = 640,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer = None,
@@ -71,6 +72,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/streaming_transducer.py

@@ -195,6 +195,7 @@ def __init__(self,
                  joint_dim: int = 640,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer = None,
@@ -226,6 +227,7 @@ def __init__(self,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,

tensorflow_asr/models/transducer.py

@@ -165,6 +165,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  kernel_regularizer=None,
                  bias_regularizer=None,
@@ -183,6 +184,7 @@ def __init__(self,
             raise ValueError("activation must be either 'linear', 'relu' or 'tanh'")
 
         self.prejoint_linear = prejoint_linear
+        self.postjoint_linear = postjoint_linear
 
         if self.prejoint_linear:
             self.ffn_enc = tf.keras.layers.Dense(
@@ -205,6 +207,13 @@ def __init__(self,
         else:
             raise ValueError("joint_mode must be either 'add' or 'concat'")
 
+        if self.postjoint_linear:
+            self.ffn = tf.keras.layers.Dense(
+                joint_dim, name=f"{name}_ffn",
+                kernel_regularizer=kernel_regularizer,
+                bias_regularizer=bias_regularizer
+            )
+
         self.ffn_out = tf.keras.layers.Dense(
             vocabulary_size, name=f"{name}_vocab",
             kernel_regularizer=kernel_regularizer,
@@ -221,6 +230,8 @@ def call(self, inputs, training=False, **kwargs):
         enc_out = self.enc_reshape(enc_out, repeats=tf.shape(pred_out)[1])
         pred_out = self.pred_reshape(pred_out, repeats=tf.shape(enc_out)[1])
         outputs = self.joint([enc_out, pred_out], training=training)
+        if self.postjoint_linear:
+            outputs = self.ffn(outputs, training=training)
         outputs = self.activation(outputs, training=training)  # => [B, T, U, V]
         outputs = self.ffn_out(outputs, training=training)
         return outputs
@@ -231,7 +242,7 @@ def get_config(self):
         conf.update(self.ffn_out.get_config())
         conf.update(self.activation.get_config())
         conf.update(self.joint.get_config())
-        conf.update({"prejoint_linear": self.prejoint_linear})
+        conf.update({"prejoint_linear": self.prejoint_linear, "postjoint_linear": self.postjoint_linear})
         return conf
 
 
@@ -253,6 +264,7 @@ def __init__(self,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
+                 postjoint_linear: bool = False,
                  joint_mode: str = "add",
                  joint_trainable: bool = True,
                  kernel_regularizer=None,
@@ -281,6 +293,7 @@ def __init__(self,
             joint_dim=joint_dim,
             activation=joint_activation,
             prejoint_linear=prejoint_linear,
+            postjoint_linear=postjoint_linear,
             joint_mode=joint_mode,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,