Repository: OpenNMT/OpenNMT-tf
Pull request #959 (2022-07-28 10:48:42)

Avoid building batches with more tokens than the user batch size


Patch:
diff --git a/opennmt/data/dataset.py b/opennmt/data/dataset.py
index 0d7c6e1f..ef193809 100644
--- a/opennmt/data/dataset.py
+++ b/opennmt/data/dataset.py
@@ -415,49 +415,53 @@ def batch_sequence_dataset(
     """
     batch_size = batch_size * batch_multiplier
 
-    def _get_bucket_id(features, length_fn):
-        default_id = tf.constant(0, dtype=tf.int32)
+    def _key_func(*args):
+        default_bucket_id = tf.constant(0, dtype=tf.int64)
         if length_fn is None:
-            return default_id
-        lengths = length_fn(features)
-        if lengths is None:
-            return default_id
-        if not isinstance(lengths, list):
-            lengths = [lengths]  # Fallback to the general case of parallel inputs.
-        lengths = [length // length_bucket_width for length in lengths]
-        return tf.reduce_max(lengths)
+            return default_bucket_id
 
-    def _key_func(*args):
         length_fns = length_fn
-        if length_fns is None:
-            length_fns = [None for _ in args]
-        elif not isinstance(length_fns, (list, tuple)):
+        if not isinstance(length_fns, (list, tuple)):
             length_fns = [length_fns]
+
         if len(length_fns) != len(args):
             raise ValueError(
                 "%d length functions were passed but this dataset contains "
                 "%d parallel elements" % (len(length_fns), len(args))
             )
-        # Take the highest bucket id.
-        bucket_id = tf.reduce_max(
+
+        # Get length of all parallel inputs.
+        all_lengths = tf.nest.flatten(
             [
-                _get_bucket_id(features, length_fn)
+                length_fn(features) if length_fn is not None else None
                 for features, length_fn in zip(args, length_fns)
             ]
         )
-        return tf.cast(bucket_id, tf.int64)
+
+        # Remove undefined lengths.
+        all_lengths = [length for length in all_lengths if length is not None]
+        if not all_lengths:
+            return default_bucket_id
+
+        # Define the bucket id based on the maximum length.
+        maximum_length = tf.reduce_max(all_lengths)
+
+        bucket_id = tf.math.maximum(
+            tf.cast(tf.math.ceil(maximum_length / length_bucket_width) - 1, tf.int64),
+            default_bucket_id,
+        )
+
+        return bucket_id
 
     def _reduce_func(unused_key, dataset):
         return dataset.apply(batch_dataset(batch_size, padded_shapes=padded_shapes))
 
     def _window_size_func(key):
-        if length_bucket_width > 1:
-            key += 1  # For length_bucket_width == 1, key 0 is unassigned.
-        size = batch_size // (key * length_bucket_width)
-        required_multiple = batch_multiplier * batch_size_multiple
-        if required_multiple > 1:
-            size = size + required_multiple - size % required_multiple
-        return tf.cast(tf.maximum(size, required_multiple), tf.int64)
+        bucket_max_length = (key + 1) * length_bucket_width
+        size = batch_size // bucket_max_length
+        # Round to the closest smaller or equal multiple.
+        size = (size // batch_size_multiple) * batch_size_multiple
+        return tf.math.maximum(size, batch_size_multiple)
 
     if length_bucket_width is None:
         if batch_type == "tokens":
diff --git a/opennmt/runner.py b/opennmt/runner.py
index 6993286e..a6d0e501 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -26,7 +26,7 @@ _CONFIG_FALLBACK = {
     "params": {},
     "train": {
         "batch_type": "examples",
-        "length_bucket_width": 1,
+        "length_bucket_width": 2,
         "sample_buffer_size": 500000,
         "save_summary_steps": 100,
     },
diff --git a/opennmt/utils/checkpoint.py b/opennmt/utils/checkpoint.py
index 19817a5a..bae74efb 100644
--- a/opennmt/utils/checkpoint.py
+++ b/opennmt/utils/checkpoint.py
@@ -225,7 +225,8 @@ def average_checkpoints(
     new_checkpoint_manager = tf.train.CheckpointManager(
         checkpoint, output_dir, max_to_keep=None
     )
-    new_checkpoint_manager.save(checkpoint_number=last_step)
+    path = new_checkpoint_manager.save(checkpoint_number=last_step)
+    tf.get_logger().info("Saved averaged checkpoint to %s", path)
     return output_dir
 
 
diff --git a/setup.py b/setup.py
index 3d2aee38..d5698eff 100644
--- a/setup.py
+++ b/setup.py
@@ -73,7 +73,7 @@ setup(
         "pyonmttok>=1.29.0,<2",
         "pyyaml>=5.3,<7",
         "rouge>=1.0,<2",
-        "sacrebleu>=1.5.0,<2.1",
+        "sacrebleu>=1.5.0,<2.3",
         "tensorflow-addons>=0.16,<0.18",
     ],
     extras_require={


Test patch:
diff --git a/opennmt/tests/dataset_test.py b/opennmt/tests/dataset_test.py
index 295e2ce9..3c217bf8 100644
--- a/opennmt/tests/dataset_test.py
+++ b/opennmt/tests/dataset_test.py
@@ -246,11 +246,10 @@ class DatasetTest(tf.test.TestCase):
             self.assertAllEqual(element["length"], length)
             self.assertAllEqual(element["index"], index)
 
-        self.assertEqual(len(elements), 4)
-        _check_element(elements[0], [8, 6, 7], [0, 3, 4])
-        _check_element(elements[1], [2, 1], [1, 5])
-        _check_element(elements[2], [5, 3], [2, 6])
-        _check_element(elements[3], [9], [7])
+        self.assertEqual(len(elements), 3)
+        _check_element(elements[0], [2, 1, 3], [1, 5, 6])
+        _check_element(elements[1], [8, 7, 9], [0, 4, 7])
+        _check_element(elements[2], [5, 6], [2, 3])
 
 
 if __name__ == "__main__":