diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a54f4b279..d5bad7572 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/README.md b/README.md index c6aeef322..ec6e9af3c 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.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`: diff --git a/docs/installation.md b/docs/installation.md index 6983e3fab..b4b0b829e 100644 --- a/docs/installation.md +++ b/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. diff --git a/opennmt/models/model.py b/opennmt/models/model.py index 91821e60f..b2d44ed8a 100644 --- a/opennmt/models/model.py +++ b/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 diff --git a/opennmt/optimizers/utils.py b/opennmt/optimizers/utils.py index 2db8b5cab..6fcae4fa8 100644 --- a/opennmt/optimizers/utils.py +++ b/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. diff --git a/opennmt/tests/checkpoint_test.py b/opennmt/tests/checkpoint_test.py index 9d5437c42..0342d5b5c 100644 --- a/opennmt/tests/checkpoint_test.py +++ b/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(): diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py index 0405b9fb0..f4fafed3b 100644 --- a/opennmt/tests/model_test.py +++ b/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( diff --git a/opennmt/tests/optimizer_test.py b/opennmt/tests/optimizer_test.py index 739e52246..2808a557a 100644 --- a/opennmt/tests/optimizer_test.py +++ b/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) diff --git a/opennmt/tests/training_test.py b/opennmt/tests/training_test.py index 4080b01e0..b64df22a5 100644 --- a/opennmt/tests/training_test.py +++ b/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): diff --git a/opennmt/tests/vocab_test.py b/opennmt/tests/vocab_test.py index 194aa2632..73ea1b26b 100644 --- a/opennmt/tests/vocab_test.py +++ b/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) diff --git a/opennmt/training.py b/opennmt/training.py index 8b835f863..0ae43e89e 100644 --- a/opennmt/training.py +++ b/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 diff --git a/opennmt/version.py b/opennmt/version.py index fad2453f8..a4e5dd37a 100644 --- a/opennmt/version.py +++ b/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(): diff --git a/setup.py b/setup.py index 886f0cd89..239204e31 100644 --- a/setup.py +++ b/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": [