Repository: OpenNMT/OpenNMT-tf
Pull request #897 (2021-10-13 10:53:51)

Remove support for TensorFlow 2.3


Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 744a88f9..31758b6f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -44,7 +44,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        tensorflow: [2.3, 2.4, 2.5, 2.6]
+        tensorflow: [2.4, 2.5, 2.6]
 
     steps:
       - uses: actions/checkout@v2
@@ -72,8 +72,7 @@ jobs:
 
       - name: Upload coverage report
         if: matrix.tensorflow == '2.6'
-        run: |
-          bash <(curl -s https://codecov.io/bash)
+        uses: codecov/codecov-action@v2
 
 
   build-and-deploy-package:
diff --git a/README.md b/README.md
index 6a314170..858b85e8 100644
--- a/README.md
+++ b/README.md
@@ -97,7 +97,7 @@ OpenNMT-tf also implements most of the techniques commonly used to train and eva
 OpenNMT-tf requires:
 
 * Python 3.6 or above
-* TensorFlow 2.3, 2.4, 2.5, or 2.6
+* TensorFlow 2.4, 2.5, or 2.6
 
 We recommend installing it with `pip`:
 
diff --git a/docs/installation.md b/docs/installation.md
index b4f74cf1..4e30bf5a 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -5,7 +5,7 @@
 OpenNMT-tf requires:
 
 * Python 3.6 or above
-* TensorFlow 2.3, 2.4, 2.5, or 2.6
+* TensorFlow 2.4, 2.5, or 2.6
 
 For GPU support, please read the [TensorFlow documentation](https://www.tensorflow.org/install/gpu) for additional software and hardware requirements.
 
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
index 36d625dc..e7f6e572 100644
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -167,22 +167,48 @@ def load_pretrained_embeddings(
     return pretrained
 
 
-def add_sequence_controls(ids, length, start_id=None, end_id=None):
+def add_sequence_controls(
+    ids,
+    length=None,
+    start_id=None,
+    end_id=None,
+    pad_id=constants.PADDING_ID,
+):
     """Adds sequence control tokens.
 
     Args:
-      ids: Sequence of ids as 1D or 2D (batch) tensor.
-      length: Sequence length as 0D or 1D (batch) tensor.
-      start_id: Id to prepend to the sequence (set ``None`` to disable).
-      end_id: Id to append to the sequence (set ``None`` to disable).
+      ids: Sequence of token ids as a 1D ``tf.Tensor``, a 2D ``tf.Tensor``, or a
+        2D ``tf.RaggedTensor``.
+      length: Length of each sequence in the batch as a 1D ``tf.Tensor``.
+        Required if :obj:`ids` is a 2D ``tf.Tensor``.
+      start_id: Token id to prepend to the sequences (set ``None`` to disable).
+      end_id: Token id to append to the sequences (set ``None`` to disable).
+      pad_id: Padding token ID.
 
     Returns:
-      A tuple ``(ids, length)``.
+      The updated token ids, and optionally the updated length if it was initially
+      passed as argument.
+
+    Raises:
+      ValueError: if :obj:`ids` is a dense 2D ``tf.Tensor``, but :obj:`length` is not set.
+      ValueError: if the rank of :obj:`ids` is not 1 or 2.
     """
+    input_is_ragged = isinstance(ids, tf.RaggedTensor)
+    if start_id is None and end_id is None:
+        return ids if input_is_ragged or length is None else (ids, length)
+
+    batch_size = None
     rank = ids.shape.rank
-    if rank not in (1, 2):
+    if rank == 2:
+        if not input_is_ragged:
+            if length is None:
+                raise ValueError(
+                    "length argument is required when passing a dense batch of sequences"
+                )
+            ids = tf.RaggedTensor.from_tensor(ids, lengths=length)
+        batch_size = ids.nrows()
+    elif rank != 1:
         raise ValueError("Unsupported rank %d (expected 1 or 2)" % rank)
-    batch_size = tf.shape(ids)[0] if rank == 2 else None
 
     def _make_column(value):
         value = tf.constant(value, dtype=ids.dtype)
@@ -190,22 +216,23 @@ def add_sequence_controls(ids, length, start_id=None, end_id=None):
             value = tf.fill([batch_size], value)
         return tf.expand_dims(value, -1)
 
+    concat_inputs = []
     if start_id is not None:
-        start_ids = _make_column(start_id)
-        ids = tf.concat([start_ids, ids], axis=-1)
-        length += 1
-
+        concat_inputs.append(_make_column(start_id))
+    concat_inputs.append(ids)
     if end_id is not None:
-        end_ids = _make_column(end_id)
-        if batch_size is not None:
-            # Run concat on RaggedTensor to handle sequences with variable length.
-            ids = tf.RaggedTensor.from_tensor(ids, lengths=length)
-        ids = tf.concat([ids, end_ids], axis=-1)
+        concat_inputs.append(_make_column(end_id))
+
+    ids = tf.concat(concat_inputs, axis=-1)
+
+    if not input_is_ragged:
         if batch_size is not None:
-            ids = ids.to_tensor()
-        length += 1
+            ids = ids.to_tensor(default_value=pad_id)
+        if length is not None:
+            length += len(concat_inputs) - 1
+            return ids, length
 
-    return ids, length
+    return ids
 
 
 class TextInputter(Inputter):
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index a69f333d..aeea295a 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -5,7 +5,7 @@ import abc
 import tensorflow as tf
 
 from opennmt import optimizers, schedules
-from opennmt.utils import compat, exporters, losses, misc
+from opennmt.utils import exporters, losses, misc
 
 
 class Model(tf.keras.layers.Layer):
@@ -213,11 +213,6 @@ class Model(tf.keras.layers.Layer):
           - The loss.
           - The gradients.
         """
-        # TODO: clean mixed precision API when TensorFlow requirement is updated to >=2.4.
-        loss_scale_optimizer = compat.tf_any(
-            "keras.mixed_precision.LossScaleOptimizer",
-            "keras.mixed_precision.experimental.LossScaleOptimizer",
-        )
 
         def _compute_loss():
             train_loss, report_loss = self.compute_training_loss(
@@ -233,10 +228,10 @@ class Model(tf.keras.layers.Layer):
         if tf.executing_eagerly():
             with tf.GradientTape() as tape:
                 train_loss, report_loss = _compute_loss()
-                if isinstance(optimizer, loss_scale_optimizer):
+                if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
                     train_loss = optimizer.get_scaled_loss(train_loss)
             gradients = tape.gradient(train_loss, self.trainable_weights)
-            if isinstance(optimizer, loss_scale_optimizer):
+            if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
                 gradients = optimizer.get_unscaled_gradients(gradients)
 
         else:
diff --git a/opennmt/training.py b/opennmt/training.py
index f0db4c6b..ad7c324d 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -9,21 +9,7 @@ import tensorflow as tf
 
 from opennmt.inputters import text_inputter
 from opennmt.optimizers import utils as optimizer_util
-from opennmt.utils import compat, misc
-
-
-def _add_mixed_precision_wrapper(optimizer):
-    # TODO: clean mixed precision API when TensorFlow requirement is updated to >=2.4.
-    wrapper_class = None
-    wrapper_kwargs = {}
-    if compat.tf_supports("keras.mixed_precision.LossScaleOptimizer"):
-        wrapper_class = tf.keras.mixed_precision.LossScaleOptimizer
-    else:
-        wrapper_class = tf.keras.mixed_precision.experimental.LossScaleOptimizer
-        wrapper_kwargs = dict(loss_scale="dynamic")
-    if not isinstance(optimizer, wrapper_class):
-        optimizer = wrapper_class(optimizer, **wrapper_kwargs)
-    return optimizer
+from opennmt.utils import misc
 
 
 class Trainer:
@@ -50,8 +36,10 @@ class Trainer:
 
         if optimizer is None:
             raise ValueError("No optimizer is defined")
-        if self._mixed_precision:
-            optimizer = _add_mixed_precision_wrapper(optimizer)
+        if self._mixed_precision and not isinstance(
+            optimizer, tf.keras.mixed_precision.LossScaleOptimizer
+        ):
+            optimizer = tf.keras.mixed_precision.LossScaleOptimizer(optimizer)
         self._optimizer = optimizer
 
     @property
@@ -430,15 +418,7 @@ class MirroredStrategyTrainer(Trainer):
         # We prefer not to use experimental_distribute_dataset here because it
         # sometimes fails to split the batches (noticed with tokens batch type).
         dataset_fn = dataset if callable(dataset) else lambda _: dataset
-        # TODO: clean this API usage when TensorFlow requirement is updated to >=2.4.
-        distribute_fn = getattr(
-            self._strategy, "distribute_datasets_from_function", None
-        )
-        if distribute_fn is None:
-            distribute_fn = (
-                self._strategy.experimental_distribute_datasets_from_function
-            )
-        return distribute_fn(dataset_fn)
+        return self._strategy.distribute_datasets_from_function(dataset_fn)
 
     def _forward(self, source, target, accum_steps=1, report_steps=None):
         per_replica_loss = self._strategy.run(
diff --git a/opennmt/utils/__init__.py b/opennmt/utils/__init__.py
index c95da9a1..4d03bca6 100644
--- a/opennmt/utils/__init__.py
+++ b/opennmt/utils/__init__.py
@@ -20,6 +20,7 @@ from opennmt.utils.exporters import (
     CTranslate2Exporter,
     CTranslate2Float16Exporter,
     CTranslate2Int8Exporter,
+    CTranslate2Int8Float16Exporter,
     CTranslate2Int16Exporter,
     Exporter,
     SavedModelExporter,
diff --git a/opennmt/utils/exporters.py b/opennmt/utils/exporters.py
index 6bdc1702..47b75460 100644
--- a/opennmt/utils/exporters.py
+++ b/opennmt/utils/exporters.py
@@ -157,7 +157,7 @@ class CTranslate2Exporter(Exporter):
         # Fail now if ctranslate2 package is missing.
         import ctranslate2
 
-        accepted_quantization = ("int8", "int16", "float16")
+        accepted_quantization = ("int8", "int16", "float16", "int8_float16")
         if quantization is not None and quantization not in accepted_quantization:
             raise ValueError(
                 "Invalid quantization '%s' for CTranslate2, accepted values are: %s"
@@ -211,3 +211,11 @@ class CTranslate2Float16Exporter(CTranslate2Exporter):
 
     def __init__(self):
         super().__init__(quantization="float16")
+
+
+@register_exporter(name="ctranslate2_int8_float16")
+class CTranslate2Int8Float16Exporter(CTranslate2Exporter):
+    """CTranslate2 exporter with int8_float16 quantization."""
+
+    def __init__(self):
+        super().__init__(quantization="int8_float16")
diff --git a/opennmt/utils/misc.py b/opennmt/utils/misc.py
index 4fa96389..3bbfe0d9 100644
--- a/opennmt/utils/misc.py
+++ b/opennmt/utils/misc.py
@@ -13,8 +13,6 @@ import tensorflow as tf
 
 from tensorflow.python.training.tracking import graph_view
 
-from opennmt.utils import compat
-
 
 def get_devices(count=1, fallback_to_cpu=True):
     """Gets devices.
@@ -49,17 +47,6 @@ def get_devices(count=1, fallback_to_cpu=True):
     return devices[0:count]
 
 
-# TODO: clean mixed precision API when TensorFlow requirement is updated to >=2.4.
-_set_global_policy = compat.tf_any(
-    "keras.mixed_precision.set_global_policy",
-    "keras.mixed_precision.experimental.set_policy",
-)
-_get_global_policy = compat.tf_any(
-    "keras.mixed_precision.global_policy",
-    "keras.mixed_precision.experimental.global_policy",
-)
-
-
 def enable_mixed_precision(force=False):
     """Globally enables mixed precision if the detected hardware supports it.
 
@@ -91,18 +78,18 @@ def enable_mixed_precision(force=False):
             )
             return False
 
-    _set_global_policy("mixed_float16")
+    tf.keras.mixed_precision.set_global_policy("mixed_float16")
     return True
 
 
 def disable_mixed_precision():
     """Globally disables mixed precision."""
-    _set_global_policy("float32")
+    tf.keras.mixed_precision.set_global_policy("float32")
 
 
 def mixed_precision_enabled():
     """Returns ``True`` if mixed precision is enabled."""
-    policy = _get_global_policy()
+    policy = tf.keras.mixed_precision.global_policy()
     return "float16" in policy.name
 
 
diff --git a/opennmt/version.py b/opennmt/version.py
index c530377c..29aa91b2 100644
--- a/opennmt/version.py
+++ b/opennmt/version.py
@@ -2,7 +2,7 @@
 
 __version__ = "2.22.0"
 
-INCLUSIVE_MIN_TF_VERSION = "2.3.0"
+INCLUSIVE_MIN_TF_VERSION = "2.4.0"
 EXCLUSIVE_MAX_TF_VERSION = "2.7.0"
 
 
diff --git a/setup.py b/setup.py
index cc7f8d54..aad6c0ba 100644
--- a/setup.py
+++ b/setup.py
@@ -64,7 +64,7 @@ setup(
     keywords="tensorflow opennmt nmt neural machine translation",
     python_requires=">=3.6",
     install_requires=[
-        "ctranslate2>=2.0.0,<3;platform_system=='Linux' or platform_system=='Darwin'",
+        "ctranslate2>=2.3.2,<3;platform_system=='Linux' or platform_system=='Darwin'",
         "pyonmttok>=1.29.0,<2",
         "pyyaml>=5.3,<5.5",
         "rouge>=1.0,<2",


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index 3a5add3e..0a0ee96b 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -138,7 +138,7 @@ class InputterTest(tf.test.TestCase):
     @parameterized.expand(
         [
             [[3, 4], 2, 1, 2, [1, 3, 4, 2], 4],
-            [[[3, 4], [5, 6]], [2, 1], 1, None, [[1, 3, 4], [1, 5, 6]], [3, 2]],
+            [[[3, 4], [5, 6]], [2, 1], 1, None, [[1, 3, 4], [1, 5, 0]], [3, 2]],
             [[[3, 4], [5, 6]], [2, 1], None, 2, [[3, 4, 2], [5, 2, 0]], [3, 2]],
         ]
     )
@@ -153,6 +153,11 @@ class InputterTest(tf.test.TestCase):
         self.assertAllEqual(self.evaluate(ids), expected_ids)
         self.assertAllEqual(self.evaluate(length), expected_length)
 
+    def testAddSequenceControlsRagged(self):
+        ids = tf.RaggedTensor.from_tensor([[4, 5, 6], [3, 4, 0]], padding=0)
+        ids = inputters.add_sequence_controls(ids, start_id=1, end_id=2)
+        self.assertAllEqual(ids.to_list(), [[1, 4, 5, 6, 2], [1, 3, 4, 2]])
+
     def _checkFeatures(self, features, expected_shapes):
         for name, expected_shape in expected_shapes.items():
             self.assertIn(name, features)
diff --git a/opennmt/tests/misc_test.py b/opennmt/tests/misc_test.py
index 4e14ca6a..3416f542 100644
--- a/opennmt/tests/misc_test.py
+++ b/opennmt/tests/misc_test.py
@@ -1,3 +1,5 @@
+import itertools
+
 import numpy as np
 import tensorflow as tf
 
@@ -8,15 +10,12 @@ from opennmt.utils import misc
 
 
 class MiscTest(tf.test.TestCase):
-    @parameterized.expand(
-        [
-            (False, "float32"),
-            (True, "float32"),
-            (False, "mixed_float16"),
-        ]
-    )
-    def testGetVariableName(self, distributed_variables, dtype_policy):
-        tf.keras.mixed_precision.experimental.set_policy(dtype_policy)
+    @parameterized.expand(itertools.product((True, False), repeat=2))
+    def testGetVariableName(self, distributed_variables, mixed_precision):
+        if mixed_precision:
+            misc.enable_mixed_precision(force=True)
+            self.assertTrue(misc.mixed_precision_enabled())
+
         if distributed_variables:
             devices = tf.config.list_logical_devices(device_type="CPU")
             strategy = tf.distribute.MirroredStrategy(devices=devices)
@@ -46,7 +45,10 @@ class MiscTest(tf.test.TestCase):
 
         variables = misc.get_variables_name_mapping(model, root_key="model")
         self.assertIs(variables[expected_name], variable)
-        tf.keras.mixed_precision.experimental.set_policy("float32")
+
+        if mixed_precision:
+            misc.disable_mixed_precision()
+            self.assertFalse(misc.mixed_precision_enabled())
 
     def testSetDropout(self):
         class Layer(tf.keras.layers.Layer):
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
index a0e027ce..8f7cbf3a 100644
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -3,8 +3,6 @@ import os
 import shutil
 import unittest
 
-from distutils.version import LooseVersion
-
 import tensorflow as tf
 
 from parameterized import parameterized, parameterized_class
@@ -150,11 +148,6 @@ class RunnerTest(tf.test.TestCase):
 
     @test_util.run_with_mixed_precision
     def testTrainMixedPrecision(self):
-        if tf.config.functions_run_eagerly() and LooseVersion(tf.__version__) < "2.4.0":
-            # TODO: remove this skipTest once TensorFlow requirement is updated to >=2.4.
-            self.skipTest(
-                "Mixed precision not working with TensorFlow 2.3 + eager execution"
-            )
         self.assertTrue(misc.mixed_precision_enabled())
         ar_file, en_file = self._makeTransliterationData()
         config = {