Implementation of layer-norm in the training script

.gitmodules

@@ -4,4 +4,4 @@
 	branch = master
 [submodule "tensorflow"]
 	path = tensorflow
-	url = https://github.com/mozilla/tensorflow.git
+	url = https://github.com/bernardohenz/tensorflow.git

taskcluster/.build.yml

@@ -26,7 +26,7 @@ build:
   nc_asset_name: 'native_client.tar.xz'
   args:
     tests_cmdline: ''
-  tensorflow_git_desc: 'TensorFlow: v2.3.0-4-g4336a5b'
+  tensorflow_git_desc: 'TensorFlow: v2.3.0-5-g6dc2a1b'
   test_model_task: ''
   homebrew:
     url: ''

taskcluster/.shared.yml

@@ -140,38 +140,38 @@ system:
       namespace: "project.mozilla-voice-stt.swig.win.amd64.1a4c14945012f1282c2eddc174fb7674d5295de8.0"
   tensorflow:
     linux_amd64_cpu:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.cpu/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.cpu"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.cpu/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.cpu"
     linux_amd64_cuda:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.cuda/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.cuda"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.cuda/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.cuda"
     linux_armv7:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.arm/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.arm"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.arm/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.arm"
     linux_arm64:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.arm64/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.arm64"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.arm64/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.arm64"
     darwin_amd64:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.1.osx/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.1.osx"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.osx/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.osx"
     android_arm64:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.android-arm64/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.android-arm64"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.android-arm64/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.android-arm64"
     android_armv7:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.android-armv7/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.android-armv7"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.android-armv7/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.android-armv7"
     win_amd64_cpu:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.win/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.win"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.win/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.win"
     win_amd64_cuda:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.win-cuda/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.win-cuda"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.win-cuda/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.win-cuda"
     ios_arm64:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.ios_arm64/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.ios_arm64"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.ios_arm64/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.ios_arm64"
     ios_x86_64:
-      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.ios_x86_64/artifacts/public/home.tar.xz"
-      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.4336a5b49fa6d650e24dbdba55bcef9581535244.0.ios_x86_64"
+      url: "https://community-tc.services.mozilla.com/api/index/v1/task/project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.ios_x86_64/artifacts/public/home.tar.xz"
+      namespace: "project.mozilla-voice-stt.tensorflow.pip.r2.3.6dc2a1becfd1316eb4d77240133a548e93dbff63.0.ios_x86_64"
   username: 'build-user'
   homedir:
     linux: '/home/build-user'

training/mozilla_voice_stt_training/train.py

@@ -74,7 +74,7 @@ def create_overlapping_windows(batch_x):
     return batch_x
 
 
-def dense(name, x, units, dropout_rate=None, relu=True):
+def dense(name, x, units, dropout_rate=None, relu=True, layer_norm=False):
     with tfv1.variable_scope(name):
         bias = variable_on_cpu('bias', [units], tfv1.zeros_initializer())
         weights = variable_on_cpu('weights', [x.shape[-1], units], tfv1.keras.initializers.VarianceScaling(scale=1.0, mode="fan_avg", distribution="uniform"))
@@ -84,6 +84,10 @@ def dense(name, x, units, dropout_rate=None, relu=True):
     if relu:
         output = tf.minimum(tf.nn.relu(output), FLAGS.relu_clip)
 
+    if layer_norm:
+        with tfv1.variable_scope(name):
+            output = tf.contrib.layers.layer_norm(output)
+
     if dropout_rate is not None:
         output = tf.nn.dropout(output, rate=dropout_rate)
 
@@ -177,9 +181,9 @@ def create_model(batch_x, seq_length, dropout, reuse=False, batch_size=None, pre
 
     # The next three blocks will pass `batch_x` through three hidden layers with
     # clipped RELU activation and dropout.
-    layers['layer_1'] = layer_1 = dense('layer_1', batch_x, Config.n_hidden_1, dropout_rate=dropout[0])
-    layers['layer_2'] = layer_2 = dense('layer_2', layer_1, Config.n_hidden_2, dropout_rate=dropout[1])
-    layers['layer_3'] = layer_3 = dense('layer_3', layer_2, Config.n_hidden_3, dropout_rate=dropout[2])
+    layers['layer_1'] = layer_1 = dense('layer_1', batch_x, Config.n_hidden_1, dropout_rate=dropout[0], layer_norm=FLAGS.layer_norm)
+    layers['layer_2'] = layer_2 = dense('layer_2', layer_1, Config.n_hidden_2, dropout_rate=dropout[1], layer_norm=FLAGS.layer_norm)
+    layers['layer_3'] = layer_3 = dense('layer_3', layer_2, Config.n_hidden_3, dropout_rate=dropout[2], layer_norm=FLAGS.layer_norm)
 
     # `layer_3` is now reshaped into `[n_steps, batch_size, 2*n_cell_dim]`,
     # as the LSTM RNN expects its input to be of shape `[max_time, batch_size, input_size]`.
@@ -196,7 +200,7 @@ def create_model(batch_x, seq_length, dropout, reuse=False, batch_size=None, pre
     layers['rnn_output_state'] = output_state
 
     # Now we feed `output` to the fifth hidden layer with clipped RELU activation
-    layers['layer_5'] = layer_5 = dense('layer_5', output, Config.n_hidden_5, dropout_rate=dropout[5])
+    layers['layer_5'] = layer_5 = dense('layer_5', output, Config.n_hidden_5, dropout_rate=dropout[5], layer_norm=FLAGS.layer_norm)
 
     # Now we apply a final linear layer creating `n_classes` dimensional vectors, the logits.
     layers['layer_6'] = layer_6 = dense('layer_6', layer_5, Config.n_hidden_6, relu=False)

training/mozilla_voice_stt_training/util/flags.py

@@ -135,6 +135,7 @@ def str_val_equals_help(name, val_desc):
     # Geometry
 
     f.DEFINE_integer('n_hidden', 2048, 'layer width to use when initialising layers')
+    f.DEFINE_boolean('layer_norm', False, 'wether to use layer-normalization after each fully-connected layer (except the last one)')
 
     # Initialization