Repository: OpenNMT/OpenNMT-tf
Pull request #627 (2020-03-06 14:11:29)

Disable dropout in frozen layers


Patch:
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index 5ef256e0..e6ffc12b 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -18,7 +18,6 @@ class Model(tf.keras.layers.Layer):
     super(Model, self).__init__()
     self.examples_inputter = examples_inputter
     self.params = {}
-    self._frozen_layers = False
 
   @property
   def unsupervised(self):
@@ -35,21 +34,6 @@ class Model(tf.keras.layers.Layer):
     """The inputter producing labels."""
     return getattr(self.examples_inputter, "labels_inputter", None)
 
-  @property
-  def trainable_weights(self):
-    if not self._frozen_layers:
-      self._frozen_layers = True
-      freeze_layers = self.params.get("freeze_layers")
-      if freeze_layers:
-        if not isinstance(freeze_layers, list):
-          freeze_layers = [freeze_layers]
-        for layer_path in freeze_layers:
-          layer = misc.index_structure(self, layer_path)
-          layer.trainable = False
-        tf.get_logger().info("%d weights are frozen by the freeze_layers parameter" % (
-            len(self.non_trainable_weights)))
-    return super(Model, self).trainable_weights
-
   @property
   def ctranslate2_spec(self):
     """The equivalent CTranslate2 model specification."""
@@ -86,6 +70,14 @@ class Model(tf.keras.layers.Layer):
     self.examples_inputter.initialize(data_config)
 
   def build(self, input_shape):
+    freeze_layers = self.params.get("freeze_layers")
+    if freeze_layers:
+      if not isinstance(freeze_layers, list):
+        freeze_layers = [freeze_layers]
+      for layer_path in freeze_layers:
+        layer = misc.index_structure(self, layer_path)
+        layer.trainable = False
+        misc.set_dropout(layer, 0)  # Disable dropout in frozen layers.
     self.examples_inputter.build(input_shape)
     self.built = True
 
diff --git a/opennmt/training.py b/opennmt/training.py
index 16418131..f7ae7cd5 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -166,6 +166,11 @@ class Trainer(abc.ABC):
     if first_call and self._is_master:
       self._model.visualize(self._checkpoint.model_dir)
       tf.get_logger().info("Number of model parameters: %d", self._model.count_params())
+      tf.get_logger().info(
+          "Number of model weights: %d (trainable = %d, non trainable = %d)",
+          len(self._model.weights),
+          len(self._model.trainable_weights),
+          len(self._model.non_trainable_weights))
     return training_loss, reported_loss
 
   def _save_checkpoint(self, step, moving_average=None):


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 109fa21e..aa4889cc 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -468,6 +468,8 @@ class ModelTest(tf.test.TestCase):
 
     _assert_layer_not_trainable(model.decoder.output_layer)
     _assert_layer_not_trainable(model.encoder.layers[0])
+    self.assertEqual(model.encoder.layers[0].ffn.output_dropout, 0)
+    self.assertEqual(model.encoder.layers[0].self_attention.output_dropout, 0)
 
   def testTransferWeightsNewVocab(self):