Support TensorFlow 2.11 (#987)

* Support TensorFlow 2.11

* Check version instead of import

* Fix AdamW test

.github/workflows/ci.yml

@@ -44,7 +44,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        tensorflow: [2.6, 2.7, 2.8, 2.9, '2.10']
+        tensorflow: [2.6, 2.7, 2.8, 2.9, '2.10', '2.11']
 
     steps:
       - uses: actions/checkout@v2
@@ -71,7 +71,7 @@ jobs:
           pytest --cov=opennmt --cov-report xml opennmt/tests
 
       - name: Upload coverage report
-        if: matrix.tensorflow == '2.10'
+        if: matrix.tensorflow == '2.11'
         uses: codecov/codecov-action@v2
 
 

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.7 or above
-* TensorFlow 2.6, 2.7, 2.8, 2.9, or 2.10
+* TensorFlow 2.6, 2.7, 2.8, 2.9, 2.10, or 2.11
 
 We recommend installing it with `pip`:
 

docs/installation.md

@@ -5,7 +5,7 @@
 OpenNMT-tf requires:
 
 * Python 3.7 or above
-* TensorFlow 2.6, 2.7, 2.8, 2.9, or 2.10
+* TensorFlow 2.6, 2.7, 2.8, 2.9, 2.10, or 2.11
 
 For GPU support, please read the [TensorFlow documentation](https://www.tensorflow.org/install/gpu) for additional software and hardware requirements.
 

opennmt/models/model.py

@@ -322,7 +322,7 @@ def get_optimizer(self):
         """Returns the optimizer for this model.
 
         Returns:
-          A ``tf.keras.optimizers.Optimizer`` instance or ``None`` if no optimizer
+          A ``tf.keras.optimizers.legacy.Optimizer`` instance or ``None`` if no optimizer
           is configured.
         """
         params = self.params

opennmt/optimizers/utils.py

@@ -5,13 +5,21 @@
 import tensorflow as tf
 import tensorflow_addons as tfa
 
+from packaging.version import Version
 from tensorflow_addons.optimizers.weight_decay_optimizers import (
     DecoupledWeightDecayExtension,
 )
 
 from opennmt.utils import misc
 
-_OPTIMIZERS_REGISTRY = misc.ClassRegistry(base_class=tf.keras.optimizers.Optimizer)
+if Version(tf.__version__) >= Version("2.11.0"):
+    tf_optimizers = tf.keras.optimizers.legacy
+else:
+    tf_optimizers = tf.keras.optimizers
+
+_OPTIMIZERS_REGISTRY = misc.ClassRegistry(
+    base_class=getattr(tf_optimizers, "Optimizer")
+)
 
 register_optimizer = _OPTIMIZERS_REGISTRY.register
 
@@ -23,14 +31,14 @@ def get_optimizer_class(name):
       name: The optimizer name.
 
     Returns:
-      A class extending ``tf.keras.optimizers.Optimizer``.
+      A class extending ``tf.keras.optimizers.legacy.Optimizer``.
 
     Raises:
       ValueError: if :obj:`name` can not be resolved to an optimizer class.
     """
     optimizer_class = None
     if optimizer_class is None:
-        optimizer_class = getattr(tf.keras.optimizers, name, None)
+        optimizer_class = getattr(tf_optimizers, name, None)
     if optimizer_class is None:
         optimizer_class = getattr(tfa.optimizers, name, None)
     if optimizer_class is None:
@@ -44,14 +52,14 @@ def make_optimizer(name, learning_rate, **kwargs):
     """Creates the optimizer.
 
     Args:
-      name: The name of the optimizer class in ``tf.keras.optimizers`` or
+      name: The name of the optimizer class in ``tf.keras.optimizers.legacy`` or
         ``tfa.optimizers`` as a string.
       learning_rate: The learning rate or learning rate schedule to use.
       **kwargs: Additional optimizer arguments. If ``weight_decay`` is set, the
         optimizer will be extended with decoupled weight decay.
 
     Returns:
-      A ``tf.keras.optimizers.Optimizer`` instance.
+      A ``tf.keras.optimizers.legacy.Optimizer`` instance.
 
     Raises:
       ValueError: if :obj:`name` can not be resolved to an optimizer class.

opennmt/tests/checkpoint_test.py

@@ -4,6 +4,7 @@
 
 from parameterized import parameterized
 
+from opennmt.optimizers.utils import make_optimizer
 from opennmt.utils import checkpoint as checkpoint_util
 
 
@@ -44,7 +45,7 @@ def testLastSavedStep(self):
     @parameterized.expand([(True,), (False,)])
     def testCheckpointAveraging(self, from_list):
         model = _DummyModel()
-        optimizer = tf.keras.optimizers.Adam()
+        optimizer = make_optimizer("Adam", 0.001)
 
         @tf.function
         def _build_model():

opennmt/tests/model_test.py

@@ -6,6 +6,7 @@
 from parameterized import parameterized
 
 from opennmt import decoders, encoders, inputters, models
+from opennmt.optimizers.utils import make_optimizer
 from opennmt.tests import test_util
 from opennmt.utils import misc
 
@@ -690,7 +691,7 @@ def _make_model(name, src_vocab, tgt_vocab, random_slots=False):
             model, _ = _seq2seq_model(
                 training=True, shared_embeddings=shared_embeddings
             )
-            optimizer = tf.keras.optimizers.Adam()
+            optimizer = make_optimizer("Adam", 0.001)
             data = {}
             data["source_vocabulary"] = test_util.make_data_file(
                 os.path.join(self.get_temp_dir(), "%s-src-vocab.txt" % name), src_vocab
@@ -835,7 +836,7 @@ def testBeamSearchWithMultiSourceEncoder(self):
 
     def testTrainModelOnBatch(self):
         _, _, data_config = self._makeToyEnDeData()
-        optimizer = tf.keras.optimizers.Adam()
+        optimizer = make_optimizer("Adam", 0.001)
         model = models.TransformerTiny()
         model.initialize(data_config)
         features = model.features_inputter.make_features(

opennmt/tests/optimizer_test.py

@@ -20,12 +20,12 @@ def testMakeAdamW(self):
         adam_w = utils.make_optimizer("AdamW", 0.002, weight_decay=0.1)
         self.assertIsInstance(adam_w, tfa.optimizers.AdamW)
         adam_w = utils.make_optimizer("Adam", 0.002, weight_decay=0.1)
-        self.assertIsInstance(adam_w, tf.keras.optimizers.Adam)
+        self.assertIsInstance(adam_w, utils.get_optimizer_class("Adam"))
         self.assertIsInstance(adam_w, DecoupledWeightDecayExtension)
 
     def testCustomOptimizerRegistration(self):
         @utils.register_optimizer
-        class MyCustomAdam(tf.keras.optimizers.Adam):
+        class MyCustomAdam(utils.get_optimizer_class("Adam")):
             pass
 
         optimizer = utils.make_optimizer("MyCustomAdam", 0.002)

opennmt/tests/training_test.py

@@ -4,6 +4,7 @@
 import tensorflow as tf
 
 from opennmt import inputters, models, training
+from opennmt.optimizers.utils import make_optimizer
 from opennmt.tests import test_util
 
 
@@ -63,7 +64,7 @@ def testMovingAverageDistributionStrategy(self):
 
     def testEmptyTrainingDataset(self):
         model = _make_seq2seq_model(self.get_temp_dir())
-        optimizer = tf.keras.optimizers.SGD(1.0)
+        optimizer = make_optimizer("SGD", 1.0)
         trainer = training.Trainer(model, optimizer)
 
         empty_file = os.path.join(self.get_temp_dir(), "train.txt")
@@ -78,7 +79,7 @@ def testEmptyTrainingDataset(self):
 
     def testTrainingStats(self):
         model = _make_seq2seq_model(self.get_temp_dir())
-        optimizer = tf.keras.optimizers.SGD(1.0)
+        optimizer = make_optimizer("SGD", 1.0)
         stats = training.TrainingStats(model, optimizer, warmup_steps=2)
 
         def _generate_example(length):

opennmt/tests/vocab_test.py

@@ -4,6 +4,7 @@
 import tensorflow as tf
 
 from opennmt.data import vocab as vocab_lib
+from opennmt.optimizers.utils import make_optimizer
 from opennmt.tests import test_util
 
 
@@ -155,7 +156,7 @@ def testVocabVariableUpdate(self):
 
 def _create_variable_and_slots(values):
     variable = tf.Variable(tf.constant(values, dtype=tf.float32))
-    optimizer = tf.keras.optimizers.Adam()
+    optimizer = make_optimizer("Adam", 0.001)
     optimizer._create_slots([variable])
     for slot in ("m", "v"):
         optimizer.get_slot(variable, slot).assign(variable)

opennmt/training.py

@@ -20,7 +20,7 @@ def __init__(self, model, optimizer, checkpoint=None):
 
         Args:
           model: A :class:`opennmt.models.Model` instance to train.
-          optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
+          optimizer: A ``tf.keras.optimizers.legacy.Optimizer`` instance.
           checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
             not set, no checkpoints will be saved.
         """
@@ -330,7 +330,7 @@ def __init__(self, model, optimizer, hvd, checkpoint=None):
 
         Args:
           model: A :class:`opennmt.models.Model` instance to train.
-          optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
+          optimizer: A ``tf.keras.optimizers.legacy.Optimizer`` instance.
           hvd: The global Horovod object.
           checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
             not set, no checkpoints will be saved.
@@ -385,7 +385,7 @@ def __init__(self, model, optimizer, checkpoint=None, devices=None):
 
         Args:
           model: A :class:`opennmt.models.Model` instance to train.
-          optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
+          optimizer: A ``tf.keras.optimizers.legacy.Optimizer`` instance.
           checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
             not set, no checkpoints will be saved.
           devices: List of device strings to use for training. If not set, all

opennmt/version.py

@@ -3,7 +3,7 @@
 __version__ = "2.29.1"
 
 INCLUSIVE_MIN_TF_VERSION = "2.6.0"
-EXCLUSIVE_MAX_TF_VERSION = "2.11.0"
+EXCLUSIVE_MAX_TF_VERSION = "2.12.0"
 
 
 def _check_tf_version():

setup.py

@@ -75,7 +75,7 @@ def get_project_version():
         "pyyaml>=5.3,<7",
         "rouge>=1.0,<2",
         "sacrebleu>=1.5.0,<2.3",
-        "tensorflow-addons>=0.16,<0.19",
+        "tensorflow-addons>=0.16,<0.20",
     ],
     extras_require={
         "tensorflow": [