Repository: OpenNMT/OpenNMT-tf
Pull request #617 (2020-02-18 09:33:40)

Allow different number of encoder/decoder layers in Transformer ctor


Patch:
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index 67a3892d..58970613 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -39,7 +39,8 @@ class Transformer(SequenceToSequence):
       target_inputter: A :class:`opennmt.inputters.Inputter` to process
         the target data. Currently, only the
         :class:`opennmt.inputters.WordEmbedder` is supported.
-      num_layers: The shared number of layers.
+      num_layers: The number of layers or a 2-tuple with the number of encoder
+        layers and decoder layers.
       num_units: The number of hidden units.
       num_heads: The number of heads in each self-attention layers.
       ffn_inner_dim: The inner dimension of the feed forward layers.
@@ -59,9 +60,13 @@ class Transformer(SequenceToSequence):
       maximum_relative_position: Maximum relative position representation
         (from https://arxiv.org/abs/1803.02155).
     """
+    if isinstance(num_layers, (list, tuple)):
+      num_encoder_layers, num_decoder_layers = num_layers
+    else:
+      num_encoder_layers, num_decoder_layers = num_layers, num_layers
     encoders = [
         SelfAttentionEncoder(
-            num_layers,
+            num_encoder_layers,
             num_units=num_units,
             num_heads=num_heads,
             ffn_inner_dim=ffn_inner_dim,
@@ -80,7 +85,7 @@ class Transformer(SequenceToSequence):
     else:
       encoder = encoders[0]
     decoder = SelfAttentionDecoder(
-        num_layers,
+        num_decoder_layers,
         num_units=num_units,
         num_heads=num_heads,
         ffn_inner_dim=ffn_inner_dim,
@@ -93,11 +98,13 @@ class Transformer(SequenceToSequence):
         maximum_relative_position=maximum_relative_position)
 
     self._num_units = num_units
-    self._num_layers = num_layers
+    self._num_encoder_layers = num_encoder_layers
+    self._num_decoder_layers = num_decoder_layers
     self._num_heads = num_heads
     self._with_relative_position = maximum_relative_position is not None
     self._is_ct2_compatible = (
         isinstance(encoder, SelfAttentionEncoder)
+        and num_encoder_layers == num_decoder_layers
         and ffn_activation == tf.nn.relu
         and ((self._with_relative_position and position_encoder_class is None)
              or (not self._with_relative_position
@@ -115,7 +122,7 @@ class Transformer(SequenceToSequence):
       return None
     import ctranslate2  # pylint: disable=import-outside-toplevel
     return ctranslate2.specs.TransformerSpec(
-        self._num_layers,
+        self._num_encoder_layers,
         self._num_heads,
         with_relative_position=self._with_relative_position)
 


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 538182e6..86ea565b 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -541,6 +541,17 @@ class ModelTest(tf.test.TestCase):
     except ImportError:
       self.skipTest("ctranslate2 module is not available")
 
+  def testTransformerWithDifferentEncoderDecoderLayers(self):
+    model = models.Transformer(
+        inputters.WordEmbedder(32),
+        inputters.WordEmbedder(32),
+        num_layers=(6, 3),
+        num_units=32,
+        num_heads=8,
+        ffn_inner_dim=64)
+    self.assertLen(model.encoder.layers, 6)
+    self.assertLen(model.decoder.layers, 3)
+
 
 if __name__ == "__main__":
   tf.test.main()