Repository: OpenNMT/OpenNMT-tf
Pull request #538 (2019-11-04 13:16:26)

Initialize the decoder when calling model.initialize()


Patch:
diff --git a/opennmt/decoders/decoder.py b/opennmt/decoders/decoder.py
index aebae192..1224e178 100644
--- a/opennmt/decoders/decoder.py
+++ b/opennmt/decoders/decoder.py
@@ -102,6 +102,11 @@ class Decoder(tf.keras.layers.Layer):
     history."""
     return False
 
+  @property
+  def initialized(self):
+    """Returns ``True`` if this decoder is initialized."""
+    return self.output_layer is not None
+
   def initialize(self, vocab_size=None, output_layer=None):
     """Initializes the decoder configuration.
 
@@ -119,6 +124,18 @@ class Decoder(tf.keras.layers.Layer):
         raise ValueError("One of vocab_size and output_layer must be set")
       self.output_layer = common.Dense(vocab_size)
 
+  def reuse_embeddings(self, embeddings):
+    """Reuses embeddings in the decoder output layer.
+
+    Args:
+      embeddings: The embeddings matrix to reuse.
+
+    Raises:
+      RuntimeError: if the decoder was not initialized.
+    """
+    self._assert_is_initialized()
+    self.output_layer.set_kernel(embeddings, transpose=True)
+
   def initial_state(self,
                     memory=None,
                     memory_sequence_length=None,
@@ -402,7 +419,7 @@ class Decoder(tf.keras.layers.Layer):
 
   def _assert_is_initialized(self):
     """Raises an expection if the decoder was not initialized."""
-    if self.output_layer is None:
+    if not self.initialized:
       raise RuntimeError("The decoder was not initialized")
 
   def _assert_memory_is_compatible(self, memory, memory_sequence_length):
diff --git a/opennmt/layers/common.py b/opennmt/layers/common.py
index 053995ad..247e6456 100644
--- a/opennmt/layers/common.py
+++ b/opennmt/layers/common.py
@@ -34,6 +34,20 @@ class Dense(tf.keras.layers.Dense):
       kwargs: Additional layers arguments.
     """
     super(Dense, self).__init__(units, **kwargs)
+    self.set_kernel(weight, transpose=transpose)
+
+  def set_kernel(self, weight, transpose=False):
+    """Use :obj:`weight` as the kernel weights matrix.
+
+    Args:
+      weight: The weight ot use.
+      transpose: Whether :obj:`weight` should be transposed or not.
+
+    Raises:
+      ValueError: if the layer is already built.
+    """
+    if self.built:
+      raise ValueError("The layer is already built")
     self.weight = weight
     self.transpose = transpose
 
diff --git a/opennmt/models/language_model.py b/opennmt/models/language_model.py
index d5d29d5f..0005d51f 100644
--- a/opennmt/models/language_model.py
+++ b/opennmt/models/language_model.py
@@ -3,7 +3,6 @@
 import tensorflow as tf
 
 from opennmt import inputters
-from opennmt import layers
 from opennmt.data import dataset as dataset_util
 from opennmt.models import model
 from opennmt.utils import decoding
@@ -43,16 +42,14 @@ class LanguageModel(model.SequenceGenerator):
         }
     })
 
+  def initialize(self, data_config, params=None):
+    super(LanguageModel, self).initialize(data_config, params=params)
+    self.decoder.initialize(vocab_size=self.examples_inputter.vocabulary_size)
+
   def build(self, input_shape):
     super(LanguageModel, self).build(input_shape)
-    vocab_size = self.examples_inputter.vocabulary_size
-    output_layer = None
     if self.reuse_embedding:
-      output_layer = layers.Dense(
-          vocab_size,
-          weight=self.examples_inputter.embedding,
-          transpose=True)
-    self.decoder.initialize(vocab_size=vocab_size, output_layer=output_layer)
+      self.decoder.reuse_embeddings(self.examples_inputter.embedding)
 
   def call(self, features, labels=None, training=None, step=None):
     outputs, predictions = None, None
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index 5ca7679c..a756ac55 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -9,7 +9,6 @@ import tensorflow_addons as tfa
 
 from opennmt import constants
 from opennmt import inputters
-from opennmt import layers
 
 from opennmt.data import noise
 from opennmt.data import text
@@ -117,6 +116,7 @@ class SequenceToSequence(model.SequenceGenerator):
 
   def initialize(self, data_config, params=None):
     super(SequenceToSequence, self).initialize(data_config, params=params)
+    self.decoder.initialize(vocab_size=self.labels_inputter.vocabulary_size)
     if self.params.get("contrastive_learning"):
       # Use the simplest and most effective CL_one from the paper.
       # https://www.aclweb.org/anthology/P19-1623
@@ -128,15 +128,8 @@ class SequenceToSequence(model.SequenceGenerator):
 
   def build(self, input_shape):
     super(SequenceToSequence, self).build(input_shape)
-    output_layer = None
     if EmbeddingsSharingLevel.share_target_embeddings(self.share_embeddings):
-      output_layer = layers.Dense(
-          self.labels_inputter.vocabulary_size,
-          weight=self.labels_inputter.embedding,
-          transpose=True)
-    self.decoder.initialize(
-        vocab_size=self.labels_inputter.vocabulary_size,
-        output_layer=output_layer)
+      self.decoder.reuse_embeddings(self.labels_inputter.embedding)
 
   def call(self, features, labels=None, training=None, step=None):
     # Encode the source.


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
index 8ccc4fe4..9e330044 100644
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -182,6 +182,22 @@ class ModelTest(tf.test.TestCase):
         prediction_heads=["tokens", "length", "log_probs"],
         params=params)
 
+  def testSequenceToSequenceWithSharedEmbedding(self):
+    model = models.SequenceToSequence(
+        inputters.WordEmbedder(16),
+        inputters.WordEmbedder(16),
+        encoders.SelfAttentionEncoder(2, 16, 4, 32),
+        decoders.SelfAttentionDecoder(2, 16, 4, 32),
+        share_embeddings=models.EmbeddingsSharingLevel.ALL)
+    _, _, data_config = self._makeToyEnDeData()
+    data_config["target_vocabulary"] = data_config["source_vocabulary"]
+    model.initialize(data_config)
+    self.assertTrue(model.decoder.initialized)
+    model.build(None)
+    self.assertEqual(
+        model.labels_inputter.embedding.experimental_ref(),
+        model.decoder.output_layer.weight.experimental_ref())
+
   @parameterized.expand([
       [tf.estimator.ModeKeys.EVAL],
       [tf.estimator.ModeKeys.PREDICT]])