Repository: OpenNMT/OpenNMT-tf
Pull request #978 (2022-10-07 13:04:20)

Remove recording of the global gradient norm
The reported norm was misleading because it did not take into account gradient accumulation and multi-GPU. Removing this feature for now.


Patch:
diff --git a/opennmt/data/dataset.py b/opennmt/data/dataset.py
index ef193809..84fd8a90 100644
--- a/opennmt/data/dataset.py
+++ b/opennmt/data/dataset.py
@@ -355,6 +355,88 @@ def batch_dataset(batch_size, padded_shapes=None):
     )
 
 
+def get_length_bucket_boundaries(bucket_width, maximum_length):
+    """Returns the list of length bucket boundaries.
+
+    The list includes at least one full bucket past the maximum length.
+    """
+
+    if maximum_length % bucket_width != 0:
+        maximum_length += bucket_width - maximum_length % bucket_width
+
+    last_boundary = maximum_length + bucket_width + 1
+
+    return list(range(bucket_width + 1, last_boundary + 1, bucket_width))
+
+
+def get_bucket_batch_sizes(
+    bucket_boundaries,
+    batch_size,
+    batch_type="examples",
+    batch_size_multiple=1,
+):
+    """Returns the batch size to use for each length bucket."""
+    batch_sizes = []
+
+    for boundary in bucket_boundaries:
+        if batch_type == "examples":
+            bucket_batch_size = batch_size
+        elif batch_type == "tokens":
+            maximum_length = boundary - 1
+            bucket_batch_size = max(batch_size // maximum_length, 1)
+
+        if batch_size_multiple > 1:
+            # Reduce the batch size to the previous multiple to avoid increasing the memory usage.
+            bucket_batch_size = max(
+                bucket_batch_size - bucket_batch_size % batch_size_multiple,
+                batch_size_multiple,
+            )
+
+        batch_sizes.append(bucket_batch_size)
+
+    # Batch size of 1 past the last bucket boundary.
+    batch_sizes.append(1)
+
+    return batch_sizes
+
+
+def get_element_length_func(length_fns):
+    """Returns a function to get the length for a dataset element.
+
+    If the dataset contains multiple parallel elements, the maximum length is returned.
+    """
+
+    if not isinstance(length_fns, (list, tuple)):
+        length_fns = [length_fns]
+
+    def _func(*elements):
+        if len(length_fns) != len(elements):
+            raise ValueError(
+                "%d length functions were passed but this dataset contains "
+                "%d parallel elements" % (len(length_fns), len(elements))
+            )
+
+        # Get length of all parallel inputs.
+        all_lengths = tf.nest.flatten(
+            [
+                length_fn(features) if length_fn is not None else None
+                for features, length_fn in zip(elements, length_fns)
+            ]
+        )
+
+        # Remove undefined lengths.
+        all_lengths = [length for length in all_lengths if length is not None]
+        if not all_lengths:
+            raise ValueError(
+                "The length should be defined for at least one dataset element"
+            )
+
+        # Return the maximum length of all parallel elements.
+        return tf.reduce_max(all_lengths)
+
+    return _func
+
+
 def batch_sequence_dataset(
     batch_size,
     batch_type="examples",
@@ -362,6 +444,8 @@ def batch_sequence_dataset(
     batch_size_multiple=1,
     length_bucket_width=None,
     length_fn=None,
+    maximum_length=None,
+    pad_to_bucket_boundary=False,
     padded_shapes=None,
 ):
     """Transformation that batches a dataset of sequences.
@@ -397,6 +481,10 @@ def batch_sequence_dataset(
         candidates from. ``None`` to not constrain batch formation.
       length_fn: A function or list of functions (in case of a parallel dataset)
         that take features as argument and return the associated sequence length.
+      maximum_length: If known, the maximum length or list of maximum lengths
+        (in case of a parallel dataset). This argument is required with
+        :obj:`pad_to_bucket_boundary`.
+      pad_to_bucket_boundary: Pad each batch to the length bucket boundary.
       padded_shapes: The padded shapes for this dataset. If ``None``, the shapes
         are automatically inferred from the dataset output shapes.
 
@@ -413,44 +501,60 @@ def batch_sequence_dataset(
     See Also:
       :func:`opennmt.data.batch_dataset`
     """
+    if batch_type not in ("examples", "tokens"):
+        raise ValueError(
+            "Invalid batch type: '{}'; should be 'examples' or 'tokens'".format(
+                batch_type
+            )
+        )
+
     batch_size = batch_size * batch_multiplier
 
-    def _key_func(*args):
-        default_bucket_id = tf.constant(0, dtype=tf.int64)
-        if length_fn is None:
-            return default_bucket_id
+    if length_bucket_width is None:
+        if batch_type == "tokens":
+            raise ValueError(
+                "Batch type 'tokens' requires length bucketing (the parameter "
+                "length_bucket_width should be non null)"
+            )
+        return batch_dataset(batch_size, padded_shapes=padded_shapes)
 
-        length_fns = length_fn
-        if not isinstance(length_fns, (list, tuple)):
-            length_fns = [length_fns]
+    element_length_func = get_element_length_func(length_fn)
 
-        if len(length_fns) != len(args):
+    if pad_to_bucket_boundary:
+        maximum_lengths = tf.nest.flatten(maximum_length)
+        maximum_lengths = [length for length in maximum_lengths if length is not None]
+        if not maximum_lengths:
             raise ValueError(
-                "%d length functions were passed but this dataset contains "
-                "%d parallel elements" % (len(length_fns), len(args))
+                "The maximum lengths should be known when padding batches to the length "
+                "bucket boundaries"
             )
 
-        # Get length of all parallel inputs.
-        all_lengths = tf.nest.flatten(
-            [
-                length_fn(features) if length_fn is not None else None
-                for features, length_fn in zip(args, length_fns)
-            ]
+        bucket_boundaries = get_length_bucket_boundaries(
+            length_bucket_width, max(maximum_lengths)
         )
 
-        # Remove undefined lengths.
-        all_lengths = [length for length in all_lengths if length is not None]
-        if not all_lengths:
-            return default_bucket_id
+        batch_sizes = get_bucket_batch_sizes(
+            bucket_boundaries,
+            batch_size,
+            batch_type=batch_type,
+            batch_size_multiple=batch_size_multiple,
+        )
 
-        # Define the bucket id based on the maximum length.
-        maximum_length = tf.reduce_max(all_lengths)
+        return lambda dataset: dataset.bucket_by_sequence_length(
+            element_length_func,
+            bucket_boundaries,
+            batch_sizes,
+            padded_shapes=padded_shapes,
+            pad_to_bucket_boundary=pad_to_bucket_boundary,
+            drop_remainder=True,
+        )
 
+    def _key_func(*args):
+        length = element_length_func(*args)
         bucket_id = tf.math.maximum(
-            tf.cast(tf.math.ceil(maximum_length / length_bucket_width) - 1, tf.int64),
-            default_bucket_id,
+            tf.cast(tf.math.ceil(length / length_bucket_width) - 1, tf.int64),
+            tf.constant(0, dtype=tf.int64),
         )
-
         return bucket_id
 
     def _reduce_func(unused_key, dataset):
@@ -463,25 +567,11 @@ def batch_sequence_dataset(
         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":
-            raise ValueError(
-                "Batch type 'tokens' requires length bucketing (the parameter "
-                "length_bucket_width should be non null)"
-            )
-        return batch_dataset(batch_size, padded_shapes=padded_shapes)
-
     kwargs = {}
     if batch_type == "examples":
         kwargs["window_size"] = batch_size
     elif batch_type == "tokens":
         kwargs["window_size_func"] = _window_size_func
-    else:
-        raise ValueError(
-            "Invalid batch type: '{}'; should be 'examples' or 'tokens'".format(
-                batch_type
-            )
-        )
 
     return lambda dataset: dataset.group_by_window(_key_func, _reduce_func, **kwargs)
 
@@ -494,6 +584,7 @@ def training_pipeline(
     process_fn=None,
     transform_fns=None,
     length_bucket_width=None,
+    pad_to_bucket_boundary=False,
     features_length_fn=None,
     labels_length_fn=None,
     maximum_features_length=None,
@@ -638,6 +729,11 @@ def training_pipeline(
         length_fn = [features_length_fn]
         if labels_length_fn is not None:
             length_fn.append(labels_length_fn)
+
+        maximum_length = [maximum_features_length]
+        if maximum_labels_length is not None:
+            maximum_length.append(maximum_labels_length)
+
         dataset = dataset.apply(
             batch_sequence_dataset(
                 batch_size,
@@ -646,6 +742,8 @@ def training_pipeline(
                 batch_size_multiple=batch_size_multiple,
                 length_bucket_width=length_bucket_width,
                 length_fn=length_fn,
+                maximum_length=maximum_length,
+                pad_to_bucket_boundary=pad_to_bucket_boundary,
             )
         )
         dataset = dataset.apply(filter_irregular_batches(batch_multiplier))
diff --git a/opennmt/decoders/decoder.py b/opennmt/decoders/decoder.py
index 5023764b..e87c3f9b 100644
--- a/opennmt/decoders/decoder.py
+++ b/opennmt/decoders/decoder.py
@@ -307,7 +307,6 @@ class Decoder(tf.keras.layers.Layer):
                 sampling_probability = None
             else:
                 fused_projection = False
-                tf.summary.scalar("sampling_probability", sampling_probability)
 
         batch_size, max_step, _ = misc.shape_list(inputs)
         inputs_ta = tf.TensorArray(inputs.dtype, size=max_step)
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index c7d0487d..6f48bf01 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -701,6 +701,7 @@ class ExampleInputterAdapter:
         batch_size_multiple=1,
         shuffle_buffer_size=None,
         length_bucket_width=None,
+        pad_to_bucket_boundary=False,
         maximum_features_length=None,
         maximum_labels_length=None,
         single_pass=False,
@@ -735,6 +736,7 @@ class ExampleInputterAdapter:
           length_bucket_width: The width of the length buckets to select batch
             candidates from (for efficiency). Set ``None`` to not constrain batch
             formation.
+          pad_to_bucket_boundary: Pad each batch to the length bucket boundary.
           maximum_features_length: The maximum length or list of maximum lengths of
             the features sequence(s). ``None`` to not constrain the length.
           maximum_labels_length: The maximum length of the labels sequence.
@@ -850,8 +852,11 @@ class ExampleInputterAdapter:
             batch_size_multiple=batch_size_multiple,
             transform_fns=transform_fns,
             length_bucket_width=length_bucket_width,
+            pad_to_bucket_boundary=pad_to_bucket_boundary,
             features_length_fn=features_length_fn,
             labels_length_fn=labels_length_fn,
+            maximum_features_length=maximum_features_length,
+            maximum_labels_length=maximum_labels_length,
             single_pass=single_pass,
             num_shards=num_shards,
             shard_index=shard_index,
diff --git a/opennmt/runner.py b/opennmt/runner.py
index f22439c5..6f13814b 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -227,6 +227,7 @@ class Runner:
                 batch_size_multiple=batch_size_multiple,
                 shuffle_buffer_size=train_config["sample_buffer_size"],
                 length_bucket_width=train_config["length_bucket_width"],
+                pad_to_bucket_boundary=train_config.get("pad_to_bucket_boundary"),
                 maximum_features_length=train_config.get("maximum_features_length"),
                 maximum_labels_length=train_config.get("maximum_labels_length"),
                 single_pass=train_config.get("single_pass", False),
diff --git a/opennmt/training.py b/opennmt/training.py
index 024aa181..f3cb9548 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -106,9 +106,7 @@ class Trainer:
 
             step = None
             moving_average = None
-            for i, loss in enumerate(
-                self._steps(dataset, accum_steps=accum_steps, report_steps=report_steps)
-            ):
+            for i, loss in enumerate(self._steps(dataset, accum_steps=accum_steps)):
                 if i == 0:
                     self._log_model_info()
 
@@ -206,14 +204,12 @@ class Trainer:
             dataset = dataset(tf.distribute.InputContext())
         return dataset
 
-    def _steps(self, dataset, accum_steps=1, report_steps=None):
+    def _steps(self, dataset, accum_steps=1):
         """Returns a generator over training steps (i.e. parameters update).
 
         Args:
           dataset: The training dataset.
           accum_steps: Accumulate the gradients of this many steps/batches.
-          report_steps: Report summary statistics every this many steps. This should
-            typically be used in a ``tf.summary.record_if`` context.
 
         Returns:
           A generator that yields a loss value to report for this step.
@@ -232,7 +228,6 @@ class Trainer:
             return self._forward(
                 next(iterator),
                 accum_steps=accum_steps,
-                report_steps=report_steps,
             )
 
         def _step():
@@ -279,30 +274,16 @@ class Trainer:
             len(self._model.non_trainable_weights),
         )
 
-    def _should_record_summaries(self, accum_steps, report_steps):
-        """Returns a boolean tensor to be used in tf.summary.record_if."""
-        if report_steps is None or not self.is_master:
-            return False
-        record_summaries = tf.equal(self._optimizer.iterations % report_steps, 0)
-        if accum_steps > 1:
-            record_summaries = tf.logical_and(
-                record_summaries, tf.equal(self._gradient_accumulator.step, 0)
-            )
-        return record_summaries
-
-    def _compute_gradients(self, batch, accum_steps, report_steps):
+    def _compute_gradients(self, batch, accum_steps):
         """Computes the gradient of a training example."""
-        record_summaries = self._should_record_summaries(accum_steps, report_steps)
-        with tf.summary.record_if(record_summaries):
-            features, labels = self._model.split_features_labels(batch)
-            reported_loss, gradients = self._model.compute_gradients(
-                features,
-                labels,
-                self._optimizer,
-                loss_scale=accum_steps * self.num_replicas,
-            )
-            self._training_stats.update_on_example(features, labels)
-            _summarize_gradients(gradients, record_summaries)
+        features, labels = self._model.split_features_labels(batch)
+        reported_loss, gradients = self._model.compute_gradients(
+            features,
+            labels,
+            self._optimizer,
+            loss_scale=accum_steps * self.num_replicas,
+        )
+        self._training_stats.update_on_example(features, labels)
         return reported_loss, gradients
 
     def _apply_gradients(self, gradients):
@@ -311,12 +292,11 @@ class Trainer:
             list(zip(gradients, self._model.trainable_variables))
         )
 
-    def _forward(self, batch, accum_steps=1, report_steps=None):
+    def _forward(self, batch, accum_steps=1):
         """Forwards a training example and accumulates the gradients."""
         loss, gradients = self._compute_gradients(
             batch,
             accum_steps,
-            report_steps,
         )
         if accum_steps > 1:
             self._gradient_accumulator(gradients)
@@ -418,11 +398,11 @@ class MirroredStrategyTrainer(Trainer):
         dataset_fn = dataset if callable(dataset) else lambda _: dataset
         return self._strategy.distribute_datasets_from_function(dataset_fn)
 
-    def _forward(self, batch, accum_steps=1, report_steps=None):
+    def _forward(self, batch, accum_steps=1):
         per_replica_loss = self._strategy.run(
             super()._forward,
             args=(batch,),
-            kwargs=dict(accum_steps=accum_steps, report_steps=report_steps),
+            kwargs=dict(accum_steps=accum_steps),
         )
         # TODO: this reduction could be delayed until _step is called.
         return self._strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, None)
@@ -435,20 +415,6 @@ class MirroredStrategyTrainer(Trainer):
             super()._update_moving_average(moving_average)
 
 
-def _summarize_gradients(gradients, should_record):
-    # Only compute the gradients global norm when the value is actually recorded.
-    if isinstance(should_record, bool) and not should_record:
-        return
-    tf.summary.scalar(
-        "gradients/global_norm",
-        tf.cond(
-            should_record,
-            true_fn=lambda: tf.linalg.global_norm(gradients),
-            false_fn=lambda: tf.constant(0, dtype=gradients[0].dtype),
-        ),
-    )
-
-
 class MovingAverage:
     """Object holding an exponential moving average of variables."""
 


Test patch:
diff --git a/opennmt/tests/dataset_test.py b/opennmt/tests/dataset_test.py
index 3c217bf8..924b0380 100644
--- a/opennmt/tests/dataset_test.py
+++ b/opennmt/tests/dataset_test.py
@@ -1,4 +1,5 @@
 import os
+import random
 
 import tensorflow as tf
 
@@ -251,6 +252,96 @@ class DatasetTest(tf.test.TestCase):
         _check_element(elements[1], [8, 7, 9], [0, 4, 7])
         _check_element(elements[2], [5, 6], [2, 3])
 
+    @parameterized.expand(
+        [
+            (3, 18, [4, 7, 10, 13, 16, 19, 22]),
+            (3, 19, [4, 7, 10, 13, 16, 19, 22, 25]),
+            (1, 9, [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]),
+        ]
+    )
+    def testBucketBoundaries(
+        self,
+        bucket_width,
+        maximum_length,
+        expected_boundaries,
+    ):
+        boundaries = dataset_util.get_length_bucket_boundaries(
+            bucket_width, maximum_length
+        )
+        self.assertListEqual(boundaries, expected_boundaries)
+
+    @parameterized.expand(
+        [
+            ([4, 7, 10, 13], 15, "examples", 1, [15, 15, 15, 15, 1]),
+            ([4, 7, 10, 13], 15, "examples", 8, [8, 8, 8, 8, 1]),
+            ([4, 7, 10, 13], 6, "examples", 8, [8, 8, 8, 8, 1]),
+            ([4, 7, 10, 13], 512, "tokens", 1, [170, 85, 56, 42, 1]),
+            ([4, 7, 10, 13], 512, "tokens", 8, [168, 80, 56, 40, 1]),
+        ]
+    )
+    def testBucketBatchSizes(
+        self,
+        bucket_boundaries,
+        batch_size,
+        batch_type,
+        batch_size_multiple,
+        expected_batch_sizes,
+    ):
+        batch_sizes = dataset_util.get_bucket_batch_sizes(
+            bucket_boundaries,
+            batch_size,
+            batch_type=batch_type,
+            batch_size_multiple=batch_size_multiple,
+        )
+
+        self.assertListEqual(batch_sizes, expected_batch_sizes)
+
+    def testPadToBucketBoundary(self):
+        num_examples = 1000
+        maximum_length = 100
+
+        def _random_sequence():
+            length = random.randint(1, maximum_length)
+            sequence = [1 for _ in range(length)]
+            return sequence
+
+        def _sequence_generator():
+            for _ in range(num_examples):
+                source = _random_sequence()
+                target = _random_sequence()
+                yield source, target
+
+        length_fn = lambda sequence: tf.size(sequence)
+
+        dataset = tf.data.Dataset.from_generator(
+            _sequence_generator,
+            output_signature=(
+                tf.TensorSpec([None], dtype=tf.int64),
+                tf.TensorSpec([None], dtype=tf.int64),
+            ),
+        )
+
+        dataset = dataset.apply(
+            dataset_util.batch_sequence_dataset(
+                1024,
+                "tokens",
+                batch_size_multiple=8,
+                length_bucket_width=8,
+                length_fn=[length_fn, length_fn],
+                maximum_length=[maximum_length, maximum_length],
+                pad_to_bucket_boundary=True,
+            )
+        )
+
+        elements = list(dataset)
+        self.assertNotEmpty(elements)
+
+        for source, target in elements:
+            self.assertEqual(source.shape[0] % 8, 0)
+            self.assertEqual(source.shape[1] % 8, 0)
+            self.assertEqual(target.shape[0] % 8, 0)
+            self.assertEqual(target.shape[1] % 8, 0)
+
 
 if __name__ == "__main__":
     tf.test.main()