🚀 add trainable option for fine tuning

Author: nglehuy

tensorflow_asr/models/conformer.py

@@ -378,6 +378,7 @@ def __init__(self,
                  encoder_depth_multiplier: int = 1,
                  encoder_fc_factor: float = 0.5,
                  encoder_dropout: float = 0,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 512,
                  prediction_embed_dropout: int = 0,
                  prediction_num_rnns: int = 1,
@@ -386,13 +387,15 @@ def __init__(self,
                  prediction_rnn_implementation: int = 2,
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 0,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer=L2,
                  bias_regularizer=L2,
-                 name: str = "conformer_transducer",
+                 name: str = "conformer",
                  **kwargs):
         super(Conformer, self).__init__(
             encoder=ConformerEncoder(
@@ -408,7 +411,9 @@ def __init__(self,
                 fc_factor=encoder_fc_factor,
                 dropout=encoder_dropout,
                 kernel_regularizer=kernel_regularizer,
-                bias_regularizer=bias_regularizer
+                bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
+                name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
             embed_dim=prediction_embed_dim,
@@ -419,13 +424,16 @@ def __init__(self,
             rnn_implementation=prediction_rnn_implementation,
             layer_norm=prediction_layer_norm,
             projection_units=prediction_projection_units,
+            prediction_trainable=prediction_trainable,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
-            name=name, **kwargs
+            name=name,
+            **kwargs
         )
         self.dmodel = encoder_dmodel
         self.time_reduction_factor = self.encoder.conv_subsampling.time_reduction_factor

tensorflow_asr/models/contextnet.py

@@ -197,6 +197,7 @@ def __init__(self,
                  vocabulary_size: int,
                  encoder_blocks: List[dict],
                  encoder_alpha: float = 0.5,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 512,
                  prediction_embed_dropout: int = 0,
                  prediction_num_rnns: int = 1,
@@ -205,10 +206,12 @@ def __init__(self,
                  prediction_rnn_implementation: int = 2,
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 0,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer=L2,
                  bias_regularizer=L2,
                  name: str = "contextnet",
@@ -219,6 +222,7 @@ def __init__(self,
                 alpha=encoder_alpha,
                 kernel_regularizer=kernel_regularizer,
                 bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
                 name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
@@ -229,14 +233,17 @@ def __init__(self,
             rnn_type=prediction_rnn_type,
             rnn_implementation=prediction_rnn_implementation,
             layer_norm=prediction_layer_norm,
+            prediction_trainable=prediction_trainable,
             projection_units=prediction_projection_units,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
-            name=name, **kwargs
+            name=name,
+            **kwargs
         )
         self.dmodel = self.encoder.blocks[-1].dmodel
         self.time_reduction_factor = 1

tensorflow_asr/models/keras/conformer.py

@@ -30,6 +30,7 @@ def __init__(self,
                  encoder_depth_multiplier: int = 1,
                  encoder_fc_factor: float = 0.5,
                  encoder_dropout: float = 0,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 512,
                  prediction_embed_dropout: int = 0,
                  prediction_num_rnns: int = 1,
@@ -38,13 +39,15 @@ def __init__(self,
                  prediction_rnn_implementation: int = 2,
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 0,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer=L2,
                  bias_regularizer=L2,
-                 name: str = "conformer_transducer",
+                 name: str = "conformer",
                  **kwargs):
         super(Conformer, self).__init__(
             encoder=ConformerEncoder(
@@ -60,7 +63,9 @@ def __init__(self,
                 fc_factor=encoder_fc_factor,
                 dropout=encoder_dropout,
                 kernel_regularizer=kernel_regularizer,
-                bias_regularizer=bias_regularizer
+                bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
+                name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
             embed_dim=prediction_embed_dim,
@@ -71,13 +76,16 @@ def __init__(self,
             rnn_implementation=prediction_rnn_implementation,
             layer_norm=prediction_layer_norm,
             projection_units=prediction_projection_units,
+            prediction_trainable=prediction_trainable,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
-            name=name, **kwargs
+            name=name,
+            **kwargs
         )
         self.dmodel = encoder_dmodel
         self.time_reduction_factor = self.encoder.conv_subsampling.time_reduction_factor

tensorflow_asr/models/keras/contextnet.py

@@ -25,6 +25,7 @@ def __init__(self,
                  vocabulary_size: int,
                  encoder_blocks: List[dict],
                  encoder_alpha: float = 0.5,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 512,
                  prediction_embed_dropout: int = 0,
                  prediction_num_rnns: int = 1,
@@ -33,10 +34,12 @@ def __init__(self,
                  prediction_rnn_implementation: int = 2,
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 0,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer=L2,
                  bias_regularizer=L2,
                  name: str = "contextnet",
@@ -47,6 +50,7 @@ def __init__(self,
                 alpha=encoder_alpha,
                 kernel_regularizer=kernel_regularizer,
                 bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
                 name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
@@ -58,13 +62,16 @@ def __init__(self,
             rnn_implementation=prediction_rnn_implementation,
             layer_norm=prediction_layer_norm,
             projection_units=prediction_projection_units,
+            prediction_trainable=prediction_trainable,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
-            name=name, **kwargs
+            name=name,
+            **kwargs
         )
         self.dmodel = self.encoder.blocks[-1].dmodel
         self.time_reduction_factor = 1

tensorflow_asr/models/keras/streaming_transducer.py

@@ -28,17 +28,20 @@ def __init__(self,
                  encoder_rnn_type: str = "lstm",
                  encoder_rnn_units: int = 2048,
                  encoder_layer_norm: bool = True,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 320,
                  prediction_embed_dropout: float = 0,
                  prediction_num_rnns: int = 2,
                  prediction_rnn_units: int = 2048,
                  prediction_rnn_type: str = "lstm",
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 640,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 640,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer = None,
                  bias_regularizer = None,
                  name = "StreamingTransducer",
@@ -53,6 +56,7 @@ def __init__(self,
                 layer_norm=encoder_layer_norm,
                 kernel_regularizer=kernel_regularizer,
                 bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
                 name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
@@ -63,10 +67,12 @@ def __init__(self,
             rnn_type=prediction_rnn_type,
             layer_norm=prediction_layer_norm,
             projection_units=prediction_projection_units,
+            prediction_trainable=prediction_trainable,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
             name=name, **kwargs

tensorflow_asr/models/streaming_transducer.py

@@ -183,17 +183,20 @@ def __init__(self,
                  encoder_rnn_type: str = "lstm",
                  encoder_rnn_units: int = 2048,
                  encoder_layer_norm: bool = True,
+                 encoder_trainable: bool = True,
                  prediction_embed_dim: int = 320,
                  prediction_embed_dropout: float = 0,
                  prediction_num_rnns: int = 2,
                  prediction_rnn_units: int = 2048,
                  prediction_rnn_type: str = "lstm",
                  prediction_layer_norm: bool = True,
                  prediction_projection_units: int = 640,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 640,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer = None,
                  bias_regularizer = None,
                  name = "StreamingTransducer",
@@ -208,6 +211,7 @@ def __init__(self,
                 layer_norm=encoder_layer_norm,
                 kernel_regularizer=kernel_regularizer,
                 bias_regularizer=bias_regularizer,
+                trainable=encoder_trainable,
                 name=f"{name}_encoder"
             ),
             vocabulary_size=vocabulary_size,
@@ -218,13 +222,16 @@ def __init__(self,
             rnn_type=prediction_rnn_type,
             layer_norm=prediction_layer_norm,
             projection_units=prediction_projection_units,
+            prediction_trainable=prediction_trainable,
             joint_dim=joint_dim,
             joint_activation=joint_activation,
             prejoint_linear=prejoint_linear,
             joint_mode=joint_mode,
+            joint_trainable=joint_trainable,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
-            name=name, **kwargs
+            name=name,
+            **kwargs
         )
         self.time_reduction_factor = self.encoder.time_reduction_factor
 

tensorflow_asr/models/transducer.py

@@ -249,10 +249,12 @@ def __init__(self,
                  rnn_implementation: int = 2,
                  layer_norm: bool = True,
                  projection_units: int = 0,
+                 prediction_trainable: bool = True,
                  joint_dim: int = 1024,
                  joint_activation: str = "tanh",
                  prejoint_linear: bool = True,
                  joint_mode: str = "add",
+                 joint_trainable: bool = True,
                  kernel_regularizer=None,
                  bias_regularizer=None,
                  name="transducer",
@@ -271,6 +273,7 @@ def __init__(self,
             projection_units=projection_units,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
+            trainable=prediction_trainable,
             name=f"{name}_prediction"
         )
         self.joint_net = TransducerJoint(
@@ -281,6 +284,7 @@ def __init__(self,
             joint_mode=joint_mode,
             kernel_regularizer=kernel_regularizer,
             bias_regularizer=bias_regularizer,
+            trainable=joint_trainable,
             name=f"{name}_joint"
         )
         self.time_reduction_factor = 1