Repository: OpenNMT/OpenNMT-tf
Pull request #962 (2022-07-29 08:03:28)

Pad dynamically when preparing the dataset for batch size autotuning
We want the length dimension to be unknown when building the graph.


Patch:
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index c9fa67d6..e0d26876 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -817,6 +817,25 @@ class ExampleInputterAdapter:
             padded_shapes = self.get_padded_shapes(
                 dataset.element_spec, maximum_length=maximum_length
             )
+
+            # Dynamically pad each sequence to the maximum length.
+            def _pad_to_shape(tensor, padded_shape):
+                if tensor.shape.rank == 0:
+                    return tensor
+                tensor_shape = misc.shape_list(tensor)
+                paddings = [
+                    [0, padded_dim - tensor_dim]
+                    if tf.is_tensor(tensor_dim) and padded_dim is not None
+                    else [0, 0]
+                    for tensor_dim, padded_dim in zip(tensor_shape, padded_shape)
+                ]
+                return tf.pad(tensor, paddings)
+
+            dataset = dataset.map(
+                lambda *arg: tf.nest.map_structure(
+                    _pad_to_shape, misc.item_or_tuple(arg), padded_shapes
+                )
+            )
             dataset = dataset.apply(
                 dataset_util.batch_sequence_dataset(
                     batch_size,
@@ -824,7 +843,6 @@ class ExampleInputterAdapter:
                     batch_multiplier=batch_multiplier,
                     length_bucket_width=1,
                     length_fn=constant_length_fn,
-                    padded_shapes=padded_shapes,
                 )
             )
             return dataset
diff --git a/opennmt/version.py b/opennmt/version.py
index cae77071..039ce58c 100644
--- a/opennmt/version.py
+++ b/opennmt/version.py
@@ -9,14 +9,14 @@ EXCLUSIVE_MAX_TF_VERSION = "2.10.0"
 def _check_tf_version():
     import warnings
 
-    from distutils.version import LooseVersion
-
     import tensorflow as tf
 
+    from packaging.version import Version
+
     if (
-        LooseVersion(INCLUSIVE_MIN_TF_VERSION)
-        <= LooseVersion(tf.__version__)
-        < LooseVersion(EXCLUSIVE_MAX_TF_VERSION)
+        Version(INCLUSIVE_MIN_TF_VERSION)
+        <= Version(tf.__version__)
+        < Version(EXCLUSIVE_MAX_TF_VERSION)
     ):
         return
 
diff --git a/setup.py b/setup.py
index d5698eff..e7adbed9 100644
--- a/setup.py
+++ b/setup.py
@@ -70,6 +70,7 @@ setup(
     python_requires=">=3.7",
     install_requires=[
         "ctranslate2>=2.17.0,<3",
+        "packaging",
         "pyonmttok>=1.29.0,<2",
         "pyyaml>=5.3,<7",
         "rouge>=1.0,<2",


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index f273d40c..b7bd8ae5 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -650,6 +650,9 @@ class InputterTest(tf.test.TestCase):
             batch_autotune_mode=True,
         )
 
+        source_spec, target_spec = dataset.element_spec
+        self.assertListEqual(source_spec["ids"].shape.as_list(), [None, None])
+
         source, target = next(iter(dataset))
         self.assertListEqual(source["ids"].shape.as_list(), [8, 100])
         self.assertListEqual(target["ids"].shape.as_list(), [8, 120])