Repository: OpenNMT/OpenNMT-tf
Pull request #830 (2021-05-06 07:24:48)

Revert "Allow non fully defined gradient shape when accumulating gradients"
Reverts OpenNMT/OpenNMT-tf#825

This was a regression in TensorFlow 2.5 RC that will be fixed in the final release: https://github.com/tensorflow/tensorflow/issues/48430


Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 4155b5f3..ffab6416 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -60,6 +60,7 @@ jobs:
           wget https://s3.amazonaws.com/opennmt-models/transliteration-aren-v2.tar.gz
           mkdir -p testdata
           tar xf transliteration-aren-v2.tar.gz -C testdata
+          wget -P testdata https://opennmt-trainingdata.s3.amazonaws.com/wmtende.model
 
       - name: Run tests
         run: |
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index 1b724eba..72e4de8c 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -279,7 +279,7 @@ class TextInputter(Inputter):
             return features
         if "text" in features:
             element = features.pop("text")
-        tokens = self.tokenizer.tokenize(element)
+        tokens = self.tokenizer.tokenize(element, training=training)
         if isinstance(tokens, tf.RaggedTensor):
             length = tokens.row_lengths()
             tokens = tokens.to_tensor()
diff --git a/opennmt/optimizers/utils.py b/opennmt/optimizers/utils.py
index 5b1ebb43..98c7a09f 100644
--- a/opennmt/optimizers/utils.py
+++ b/opennmt/optimizers/utils.py
@@ -136,9 +136,6 @@ class GradientAccumulator(object):
             return
         self._accum_steps.assign(0)
         for gradient in self._gradients:
-            shape = (
-                gradient.shape
-                if gradient.shape.is_fully_defined()
-                else tf.shape(gradient)
+            gradient.assign(
+                tf.zeros(gradient.shape, dtype=gradient.dtype), read_value=False
             )
-            gradient.assign(tf.zeros(shape, dtype=gradient.dtype), read_value=False)
diff --git a/opennmt/tokenizers/opennmt_tokenizer.py b/opennmt/tokenizers/opennmt_tokenizer.py
index 4762835f..0c099f5d 100644
--- a/opennmt/tokenizers/opennmt_tokenizer.py
+++ b/opennmt/tokenizers/opennmt_tokenizer.py
@@ -52,8 +52,8 @@ class OpenNMTTokenizer(tokenizer.Tokenizer):
 
         return assets
 
-    def _tokenize_string(self, text):
-        tokens, _ = self._tokenizer.tokenize(text)
+    def _tokenize_string(self, text, training):
+        tokens, _ = self._tokenizer.tokenize(text, training=training)
         return tokens
 
     def _detokenize_string(self, tokens):
diff --git a/opennmt/tokenizers/tokenizer.py b/opennmt/tokenizers/tokenizer.py
index 461ec7a8..6d918683 100644
--- a/opennmt/tokenizers/tokenizer.py
+++ b/opennmt/tokenizers/tokenizer.py
@@ -31,7 +31,11 @@ class Tokenizer(abc.ABC):
         return {}
 
     def tokenize_stream(
-        self, input_stream=sys.stdin, output_stream=sys.stdout, delimiter=" "
+        self,
+        input_stream=sys.stdin,
+        output_stream=sys.stdout,
+        delimiter=" ",
+        training=True,
     ):
         """Tokenizes a stream of sentences.
 
@@ -39,15 +43,19 @@ class Tokenizer(abc.ABC):
           input_stream: The input stream.
           output_stream: The output stream.
           delimiter: The token delimiter to use for text serialization.
+          training: Set to ``False`` to tokenize for inference.
         """
         for line in input_stream:
             line = line.strip()
-            tokens = self.tokenize(line)
+            tokens = self.tokenize(line, training=training)
             merged_tokens = delimiter.join(tokens)
             misc.print_as_bytes(merged_tokens, stream=output_stream)
 
     def detokenize_stream(
-        self, input_stream=sys.stdin, output_stream=sys.stdout, delimiter=" "
+        self,
+        input_stream=sys.stdin,
+        output_stream=sys.stdout,
+        delimiter=" ",
     ):
         """Detokenizes a stream of sentences.
 
@@ -61,12 +69,13 @@ class Tokenizer(abc.ABC):
             string = self.detokenize(tokens)
             misc.print_as_bytes(string, stream=output_stream)
 
-    def tokenize(self, text):
+    def tokenize(self, text, training=True):
         """Tokenizes text.
 
         Args:
           text: A string or batch of strings to tokenize as a ``tf.Tensor`` or
             Python values.
+          training: Set to ``False`` to tokenize for inference.
 
         Returns:
           - If :obj:`text` is a Python string, a list of Python strings.
@@ -79,22 +88,22 @@ class Tokenizer(abc.ABC):
           ValueError: if the rank of :obj:`text` is greater than 1.
         """
         with tf.device("cpu:0"):
-            return self._tokenize(text)
+            return self._tokenize(text, training)
 
-    def _tokenize(self, text):
+    def _tokenize(self, text, training):
         if tf.is_tensor(text):
             rank = len(text.shape)
             if rank == 0:
-                return self._tokenize_tensor(text)
+                return self._tokenize_tensor(text, training)
             elif rank == 1:
-                return self._tokenize_batch_tensor(text)
+                return self._tokenize_batch_tensor(text, training)
             else:
                 raise ValueError("Unsupported tensor rank %d for tokenization" % rank)
         elif isinstance(text, list):
-            return list(map(self.tokenize, text))
+            return list(map(lambda t: self._tokenize(t, training), text))
         else:
             text = tf.compat.as_text(text)
-            return self._tokenize_string(text)
+            return self._tokenize_string(text, training)
 
     def detokenize(self, tokens, sequence_length=None):
         """Detokenizes tokens.
@@ -150,7 +159,7 @@ class Tokenizer(abc.ABC):
             tokens = [tf.compat.as_text(token) for token in tokens]
             return self._detokenize_string(tokens)
 
-    def _tokenize_tensor(self, text):
+    def _tokenize_tensor(self, text, training):
         """Tokenizes a tensor.
 
         When not overriden, this default implementation calls the string-based
@@ -158,6 +167,7 @@ class Tokenizer(abc.ABC):
 
         Args:
           text: A 1-D string ``tf.Tensor``.
+          training: Set to ``False`` to tokenize for inference.
 
         Returns:
           A 1-D string ``tf.Tensor``.
@@ -165,14 +175,14 @@ class Tokenizer(abc.ABC):
 
         def _python_wrapper(string_t):
             string = tf.compat.as_text(string_t.numpy())
-            tokens = self._tokenize_string(string)
+            tokens = self._tokenize_string(string, training)
             return tf.constant(tokens, dtype=tf.string)
 
         tokens = tf.py_function(_python_wrapper, [text], tf.string)
         tokens.set_shape([None])
         return tokens
 
-    def _tokenize_batch_tensor(self, text):
+    def _tokenize_batch_tensor(self, text, training):
         """Tokenizes a batch of texts.
 
         When not overriden, this default implementation calls _tokenize_tensor on
@@ -180,6 +190,7 @@ class Tokenizer(abc.ABC):
 
         Args:
           text: A 1-D string ``tf.Tensor``.
+          training: Set to ``False`` to tokenize for inference.
 
         Returns:
           A 2-D string ``tf.RaggedTensor``.
@@ -188,7 +199,7 @@ class Tokenizer(abc.ABC):
         # spaces first and then split on spaces with a function returning a RaggedTensor.
         tokens = tf.map_fn(
             lambda x: tf.strings.reduce_join(
-                self._tokenize_tensor(x), axis=0, separator=" "
+                self._tokenize_tensor(x, training), axis=0, separator=" "
             ),
             text,
         )
@@ -249,13 +260,14 @@ class Tokenizer(abc.ABC):
         return self._detokenize_batch_tensor(tokens.to_tensor(), tokens.row_lengths())
 
     @abc.abstractmethod
-    def _tokenize_string(self, text):
+    def _tokenize_string(self, text, training):
         """Tokenizes a Python unicode string.
 
         This method should be thread-safe.
 
         Args:
           text: A Python unicode string.
+          training: Set to ``False`` to tokenize for inference.
 
         Returns:
           A list of Python unicode strings.
@@ -345,10 +357,10 @@ class SpaceTokenizer(Tokenizer):
     def in_graph(self):
         return self._in_graph
 
-    def _tokenize_tensor(self, text):
-        return self._tokenize_batch_tensor(text)
+    def _tokenize_tensor(self, text, training):
+        return self._tokenize_batch_tensor(text, training)
 
-    def _tokenize_batch_tensor(self, text):
+    def _tokenize_batch_tensor(self, text, _):
         return tf.strings.split(text)
 
     def _detokenize_tensor(self, tokens):
@@ -361,7 +373,7 @@ class SpaceTokenizer(Tokenizer):
     def _detokenize_ragged_tensor(self, tokens):
         return tf.strings.reduce_join(tokens, axis=tokens.shape.rank - 1, separator=" ")
 
-    def _tokenize_string(self, text):
+    def _tokenize_string(self, text, _):
         return text.split()
 
     def _detokenize_string(self, tokens):
@@ -376,10 +388,10 @@ class CharacterTokenizer(Tokenizer):
     def in_graph(self):
         return True
 
-    def _tokenize_tensor(self, text):
-        return self._tokenize_batch_tensor(text)
+    def _tokenize_tensor(self, text, training):
+        return self._tokenize_batch_tensor(text, training)
 
-    def _tokenize_batch_tensor(self, text):
+    def _tokenize_batch_tensor(self, text, _):
         text = tf.strings.regex_replace(text, " ", "▁")
         return tf.strings.unicode_split(text, "UTF-8")
 
@@ -394,7 +406,7 @@ class CharacterTokenizer(Tokenizer):
         text = tf.strings.reduce_join(tokens, axis=tokens.shape.rank - 1)
         return tf.strings.regex_replace(text, "▁", " ")
 
-    def _tokenize_string(self, text):
+    def _tokenize_string(self, text, _):
         return list(text.replace(" ", "▁"))
 
     def _detokenize_string(self, tokens):


Test patch:
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
index 7510e1f9..34e4483e 100644
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -17,9 +17,7 @@ from opennmt.utils import misc
 from opennmt.tests import test_util
 
 
-test_dir = os.path.dirname(os.path.realpath(__file__))
-root_dir = os.path.join(test_dir, "..", "..")
-test_data = os.path.join(root_dir, "testdata")
+test_data = test_util.get_test_data_dir()
 
 
 def _get_test_class_name(cls, num, params_dict):
diff --git a/opennmt/tests/test_util.py b/opennmt/tests/test_util.py
index 2415f098..e5e86a87 100644
--- a/opennmt/tests/test_util.py
+++ b/opennmt/tests/test_util.py
@@ -1,3 +1,4 @@
+import os
 import unittest
 
 import tensorflow as tf
@@ -15,6 +16,13 @@ def skip_if_unsupported(symbol):
     return unittest.skipIf(not compat.tf_supports(symbol), "tf.%s is not supported")
 
 
+def get_test_data_dir():
+    test_dir = os.path.dirname(os.path.realpath(__file__))
+    root_dir = os.path.join(test_dir, "..", "..")
+    test_data = os.path.join(root_dir, "testdata")
+    return test_data
+
+
 def make_data_file(path, lines):
     with open(path, "w") as data:
         for line in lines:
diff --git a/opennmt/tests/tokenizer_test.py b/opennmt/tests/tokenizer_test.py
index 1daedb69..3114b7a1 100644
--- a/opennmt/tests/tokenizer_test.py
+++ b/opennmt/tests/tokenizer_test.py
@@ -1,37 +1,42 @@
 import os
 import yaml
+import unittest
 
 import tensorflow as tf
 
 from opennmt import tokenizers
+from opennmt.tests import test_util
+
+
+sp_model = os.path.join(test_util.get_test_data_dir(), "wmtende.model")
 
 
 class TokenizerTest(tf.test.TestCase):
-    def _testTokenizerOnTensor(self, tokenizer, text, ref_tokens):
+    def _testTokenizerOnTensor(self, tokenizer, text, ref_tokens, training):
         ref_tokens = [tf.compat.as_bytes(token) for token in ref_tokens]
         text = tf.constant(text)
-        tokens = tokenizer.tokenize(text)
+        tokens = tokenizer.tokenize(text, training=training)
         self.assertIsInstance(tokens, tf.Tensor)
         self.assertAllEqual(ref_tokens, self.evaluate(tokens))
 
-    def _testTokenizerOnBatchTensor(self, tokenizer, text, ref_tokens):
+    def _testTokenizerOnBatchTensor(self, tokenizer, text, ref_tokens, training):
         text = tf.constant(text)
-        tokens = tokenizer.tokenize(text)
+        tokens = tokenizer.tokenize(text, training=training)
         self.assertIsInstance(tokens, tf.RaggedTensor)
         self.assertAllEqual(
             tokens.to_list(), tf.nest.map_structure(tf.compat.as_bytes, ref_tokens)
         )
 
-    def _testTokenizerOnString(self, tokenizer, text, ref_tokens):
-        tokens = tokenizer.tokenize(text)
+    def _testTokenizerOnString(self, tokenizer, text, ref_tokens, training):
+        tokens = tokenizer.tokenize(text, training=training)
         self.assertAllEqual(ref_tokens, tokens)
 
-    def _testTokenizer(self, tokenizer, text, ref_tokens):
-        self.assertAllEqual(tokenizer.tokenize(text), ref_tokens)
-        self._testTokenizerOnBatchTensor(tokenizer, text, ref_tokens)
+    def _testTokenizer(self, tokenizer, text, ref_tokens, training=True):
+        self.assertAllEqual(tokenizer.tokenize(text, training), ref_tokens)
+        self._testTokenizerOnBatchTensor(tokenizer, text, ref_tokens, training)
         for txt, ref in zip(text, ref_tokens):
-            self._testTokenizerOnTensor(tokenizer, txt, ref)
-            self._testTokenizerOnString(tokenizer, txt, ref)
+            self._testTokenizerOnTensor(tokenizer, txt, ref, training)
+            self._testTokenizerOnString(tokenizer, txt, ref, training)
 
     def _testDetokenizerOnTensor(self, tokenizer, tokens, ref_text):
         ref_text = tf.compat.as_bytes(ref_text)
@@ -140,6 +145,16 @@ class TokenizerTest(tf.test.TestCase):
             tokenizer, ["Hello World-s"], [["Hello", "▁", "World", "-", "s"]]
         )
 
+    @unittest.skipIf(not os.path.isfile(sp_model), "Missing SentencePiece model")
+    def testOpenNMTTokenizerInferenceMode(self):
+        tokenizer = tokenizers.OpenNMTTokenizer(
+            mode="none",
+            sp_model_path=sp_model,
+            sp_nbest_size=64,
+            sp_alpha=0.1,
+        )
+        self._testTokenizer(tokenizer, ["appealing"], [["▁appealing"]], training=False)
+
     def testOpenNMTTokenizerAssets(self):
         asset_dir = self.get_temp_dir()
         # Write a dummy BPE model.
diff --git a/opennmt/tests/transformer_test.py b/opennmt/tests/transformer_test.py
index 31d6338a..ca3a070b 100644
--- a/opennmt/tests/transformer_test.py
+++ b/opennmt/tests/transformer_test.py
@@ -159,6 +159,20 @@ class TransformerTest(tf.test.TestCase):
         cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
         _, cache = attention(x, cache=cache)
 
+    def testMultiHeadSelfAttentionRelativeGradients(self):
+        attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
+
+        @tf.function
+        def _compute_gradients_in_function(x):
+            with tf.GradientTape() as tape:
+                y, _ = attention(x)
+                loss = tf.math.reduce_sum(y)
+            gradients = tape.gradient(loss, attention.weights)
+            for gradient in gradients:
+                self.assertTrue(gradient.shape.is_fully_defined())
+
+        _compute_gradients_in_function(tf.random.uniform([4, 1, 10]))
+
     def testMultiHeadAttention(self):
         attention = transformer.MultiHeadAttention(4, 20, return_attention=True)
         queries = tf.random.uniform([4, 5, 10])