Repository: OpenNMT/OpenNMT-tf
Pull request #806 (2021-03-09 17:06:24)

Autotune batch size on maximum sequence lengths


Patch:
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
index 4daa569c..d6a7649d 100644
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -169,6 +169,26 @@ class Inputter(tf.keras.layers.Layer):
         _ = ignore_special_tokens
         return features.get("length")
 
+    def get_padded_shapes(self, element_spec, maximum_length=None):
+        """Returns the padded shapes for dataset elements.
+
+        For example, this is used during batch size autotuning to pad all batches
+        to the maximum sequence length.
+
+        Args:
+          element_spec: A nested structure of ``tf.TensorSpec``.
+          maximum_length: Pad batches to this maximum length.
+
+        Returns:
+          A nested structure of ``tf.TensorShape``.
+        """
+        return tf.nest.map_structure(
+            lambda spec: spec.shape
+            if spec.shape.rank == 0
+            else tf.TensorShape([maximum_length]).concatenate(spec.shape[1:]),
+            element_spec,
+        )
+
     @abc.abstractmethod
     def make_features(self, element=None, features=None, training=None):
         """Creates features from data.
@@ -425,6 +445,32 @@ class ParallelInputter(MultiInputter):
         else:
             return lengths[0]
 
+    def get_padded_shapes(self, element_spec, maximum_length=None):
+        if maximum_length is None:
+            maximum_length = [None for _ in self.inputters]
+        elif not isinstance(maximum_length, (list, tuple)) or len(
+            maximum_length
+        ) != len(self.inputters):
+            raise ValueError(
+                "A maximum length should be set for each parallel inputter"
+            )
+        if self.combine_features:
+            shapes = {}
+            for i, (inputter, length) in enumerate(zip(self.inputters, maximum_length)):
+                prefix = "inputter_%d_" % i
+                spec = misc.extract_prefixed_keys(element_spec, prefix)
+                sub_shapes = inputter.get_padded_shapes(spec, maximum_length=length)
+                for key, value in sub_shapes.items():
+                    shapes["%s%s" % (prefix, key)] = value
+            return shapes
+        else:
+            return type(element_spec)(
+                inputter.get_padded_shapes(spec, maximum_length=length)
+                for inputter, spec, length in zip(
+                    self.inputters, element_spec, maximum_length
+                )
+            )
+
     def make_features(self, element=None, features=None, training=None):
         if self.combine_features:
             if features is None:
@@ -656,6 +702,7 @@ class ExampleInputterAdapter:
         prefetch_buffer_size=None,
         cardinality_multiple=1,
         weights=None,
+        batch_autotune_mode=False,
     ):
         """Builds a dataset to be used for training. It supports the full training
         pipeline, including:
@@ -695,6 +742,8 @@ class ExampleInputterAdapter:
             this value when :obj:`single_pass` is ``True``.
           weights: An optional list of weights to create a weighted dataset out of
             multiple training files.
+          batch_autotune_mode: When enabled, all batches are padded to the maximum
+            sequence length.
 
         Returns:
           A ``tf.data.Dataset``.
@@ -713,6 +762,8 @@ class ExampleInputterAdapter:
             features_length_fn = self.get_length
             labels_length_fn = None
 
+        dataset = self.make_dataset(data_files, training=True)
+
         map_fn = lambda *arg: self.make_features(
             element=misc.item_or_tuple(arg), training=True
         )
@@ -726,7 +777,46 @@ class ExampleInputterAdapter:
             lambda dataset: dataset.filter(filter_fn),
         ]
 
-        dataset = self.make_dataset(data_files, training=True)
+        if batch_autotune_mode:
+            # In this mode we want to return batches where all sequences are padded
+            # to the maximum possible length in order to maximize the memory usage.
+            # Shuffling, sharding, prefetching, etc. are not applied since correctness and
+            # performance are not important.
+
+            if isinstance(dataset, list):  # Ignore weighted dataset.
+                dataset = dataset[0]
+
+            # We repeat the dataset now to ensure full batches are always returned.
+            dataset = dataset.repeat()
+            for transform_fn in transform_fns:
+                dataset = dataset.apply(transform_fn)
+
+            # length_fn returns the maximum length instead of the actual example length so
+            # that batches are built as if each example has the maximum length.
+            if labels_file is not None:
+                constant_length_fn = [
+                    lambda x: maximum_features_length,
+                    lambda x: maximum_labels_length,
+                ]
+            else:
+                constant_length_fn = lambda x: maximum_features_length
+
+            # The length dimension is set the maximum length in the padded shapes.
+            padded_shapes = self.get_padded_shapes(
+                dataset.element_spec, maximum_length=maximum_length
+            )
+            dataset = dataset.apply(
+                dataset_util.batch_sequence_dataset(
+                    batch_size,
+                    batch_type=batch_type,
+                    batch_multiplier=batch_multiplier,
+                    length_bucket_width=1,
+                    length_fn=constant_length_fn,
+                    padded_shapes=padded_shapes,
+                )
+            )
+            return dataset
+
         if weights is not None:
             dataset = (dataset, weights)
         dataset = dataset_util.training_pipeline(
diff --git a/opennmt/runner.py b/opennmt/runner.py
index 7e7b026d..c759f03f 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -6,7 +6,6 @@ import sys
 import random
 import math
 import subprocess
-import time
 import tempfile
 import yaml
 
@@ -127,17 +126,23 @@ class Runner(object):
             # Auto tune batch size.
             if batch_size is None or batch_size == 0:
                 if train_config["batch_type"] == "examples":
-                    raise ValueError(
-                        'Batch size autotuning is only supported for the "tokens" batch type'
-                    )
-                max_batch_size = 16384
+                    min_batch_size = 1
+                    max_batch_size = 512
+                    min_range = 16
+                else:
+                    min_batch_size = 256
+                    max_batch_size = 16384
+                    min_range = 256
+
                 if train_config.get("effective_batch_size") is not None:
                     max_batch_size = min(
                         max_batch_size, train_config["effective_batch_size"]
                     )
                 train_config["batch_size"] = _auto_tune_batch_size(
                     config,
+                    min_batch_size=min_batch_size,
                     max_batch_size=max_batch_size,
+                    min_range=min_range,
                     num_devices=num_devices,
                     mixed_precision=self._mixed_precision,
                 )
@@ -222,6 +227,7 @@ class Runner(object):
                 prefetch_buffer_size=train_config.get("prefetch_buffer_size"),
                 cardinality_multiple=input_context.num_replicas_in_sync,
                 weights=data_config.get("train_files_weights"),
+                batch_autotune_mode=train_config.get("batch_autotune_mode"),
             )
         )
 
@@ -476,12 +482,12 @@ def _count_batch_accum(batch_size, target_batch_size, num_replicas=1):
 
 def _auto_tune_batch_size(
     config,
-    min_batch_size=1024,
-    max_batch_size=16384,
-    min_range=256,
-    sample_iterations=10,
+    min_batch_size,
+    max_batch_size,
+    min_range,
+    sample_iterations=5,
     num_devices=1,
-    scaling_factor=0.8,
+    scaling_factor=0.9,
     mixed_precision=False,
 ):
     """Find the largest token-based batch size that can be used with this
@@ -509,55 +515,51 @@ def _auto_tune_batch_size(
     Returns:
       The autotuned batch size.
     """
-    model_dir = config["model_dir"]
-    with tempfile.TemporaryDirectory() as tmpdir:
-        config = copy.deepcopy(config)
-        config["model_dir"] = tmpdir
-        config["train"]["save_checkpoints_steps"] = None
-        config["train"]["average_last_checkpoints"] = 0
-        config["train"]["max_step"] = sample_iterations
-        config_path = os.path.join(config["model_dir"], "batch_size_autotuner.yml")
-        model_description = os.path.join(model_dir, "model_description.py")
-
-        args = [
-            sys.executable or "python",
-            "-m",
-            "opennmt.bin.main",
-            "--config",
-            config_path,
-            "--model",
-            model_description,
-            "--checkpoint_path",
-            model_dir,
-        ]
-        if mixed_precision:
-            args.extend(["--mixed_precision"])
-        args.extend(
-            [
-                "train",
-                "--num_gpus",
-                str(num_devices),
-            ]
-        )
-
-        tf.get_logger().info(
-            "Searching the largest batch size between %d and %d with a precision of %d...",
-            min_batch_size,
-            max_batch_size,
-            min_range,
-        )
-
-        while max_batch_size - min_batch_size > min_range:
-            batch_size = (max_batch_size + min_batch_size) // 2
-
-            # Update configuration with current batch size and adjusted gradients
-            # accumulation.
-            config["train"]["batch_size"] = batch_size
+    tf.get_logger().info(
+        "Searching the largest batch size between %d and %d with a precision of %d...",
+        min_batch_size,
+        max_batch_size,
+        min_range,
+    )
+
+    model_description = os.path.join(config["model_dir"], "model_description.py")
+
+    while max_batch_size - min_batch_size > min_range:
+        batch_size = (max_batch_size + min_batch_size) // 2
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            run_config = copy.deepcopy(config)
+            run_config["model_dir"] = tmpdir
+            run_config["train"]["batch_autotune_mode"] = True
+            run_config["train"]["batch_size"] = batch_size
+            run_config["train"]["save_checkpoints_steps"] = None
+            run_config["train"]["average_last_checkpoints"] = 0
+            run_config["train"]["max_step"] = sample_iterations
+
+            config_path = os.path.join(tmpdir, "batch_size_autotuner.yml")
             with tf.io.gfile.GFile(config_path, mode="w") as config_file:
-                yaml.dump(config, config_file)
+                yaml.dump(run_config, config_file)
+
+            args = [
+                sys.executable or "python",
+                "-m",
+                "opennmt.bin.main",
+                "--config",
+                config_path,
+                "--model",
+                model_description,
+            ]
+            if mixed_precision:
+                args.extend(["--mixed_precision"])
+            args.extend(
+                [
+                    "train",
+                    "--num_gpus",
+                    str(num_devices),
+                ]
+            )
 
             tf.get_logger().info("Trying training with batch size %d...", batch_size)
-            time.sleep(1)
             with open(os.devnull, "w") as devnull:
                 process = subprocess.Popen(args, stdout=devnull, stderr=devnull)
                 exit_code = process.wait()
diff --git a/opennmt/training.py b/opennmt/training.py
index 860645a1..94bb10c2 100644
--- a/opennmt/training.py
+++ b/opennmt/training.py
@@ -30,7 +30,7 @@ def _add_mixed_precision_wrapper(optimizer):
 class Trainer:
     """Base class for model trainer, implementing single-GPU training."""
 
-    def __init__(self, model, optimizer, checkpoint=None, is_master=True):
+    def __init__(self, model, optimizer, checkpoint=None):
         """Initializes the trainer.
 
         Args:
@@ -38,10 +38,8 @@ class Trainer:
           optimizer: A ``tf.keras.optimizers.Optimizer`` instance.
           checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
             not set, no checkpoints will be saved.
-          is_master: Whether this trainer instance is the master trainer.
         """
         self._checkpoint = checkpoint
-        self._is_master = is_master
         self._model = model
         if checkpoint is not None:
             self._summary_writer = tf.summary.create_file_writer(checkpoint.model_dir)
@@ -57,6 +55,11 @@ class Trainer:
             optimizer = _add_mixed_precision_wrapper(optimizer)
         self._optimizer = optimizer
 
+    @property
+    def is_master(self):
+        """Master replica."""
+        return True
+
     @property
     def num_replicas(self):
         """Number of synchronous training replicas."""
@@ -122,7 +125,7 @@ class Trainer:
                 if tf.math.is_nan(loss):
                     raise RuntimeError("Model diverged with loss = NaN.")
 
-                if moving_average_decay is not None and self._is_master:
+                if moving_average_decay is not None and self.is_master:
                     if moving_average is None:
                         moving_average = MovingAverage(
                             self._model.trainable_variables,
@@ -135,7 +138,7 @@ class Trainer:
                 reset_throughput = False
                 self._training_stats.update_on_step(step, loss)
                 if step % report_steps == 0:
-                    self._training_stats.log(self._is_master)
+                    self._training_stats.log(self.is_master)
                     reset_throughput = True
                 if step == 1 or (save_steps is not None and step % save_steps == 0):
                     self._save_checkpoint(step, moving_average=moving_average)
@@ -163,7 +166,7 @@ class Trainer:
                     "consistent with your data."
                 )
 
-            self._training_stats.log_final(self._is_master)
+            self._training_stats.log_final(self.is_master)
             summary = self._training_stats.get_global_summary()
             self._save_checkpoint(step, moving_average=moving_average)
             self._evaluate(evaluator, step, moving_average=moving_average)
@@ -172,7 +175,7 @@ class Trainer:
     def _save_checkpoint(self, step, moving_average=None):
         """Saves a checkpoint for step."""
         if (
-            not self._is_master
+            not self.is_master
             or self._checkpoint is None
             or step == self._checkpoint.last_saved_step
         ):
@@ -188,7 +191,7 @@ class Trainer:
     def _evaluate(self, evaluator, step, moving_average=None):
         """Runs evaluation for step. Returns ``True`` is early conditions are met."""
         if (
-            not self._is_master
+            not self.is_master
             or evaluator is None
             or step == evaluator.last_evaluated_step
         ):
@@ -296,7 +299,7 @@ class Trainer:
             training_loss, variables=self._model.trainable_variables
         )
         self._training_stats.update_on_example(source, target)
-        if first_call and self._is_master:
+        if first_call and self.is_master:
             if self._checkpoint is not None:
                 self._model.visualize(self._checkpoint.model_dir)
             tf.get_logger().info(
@@ -312,7 +315,7 @@ class Trainer:
 
     def _should_record_summaries(self, report_steps, with_accum=False):
         """Returns a boolean tensor to be used in tf.summary.record_if."""
-        if report_steps is None or not self._is_master:
+        if report_steps is None or not self.is_master:
             return False
         record_summaries = tf.equal(self._optimizer.iterations % report_steps, 0)
         if with_accum:
@@ -405,11 +408,13 @@ class HorovodTrainer(Trainer):
           checkpoint: A :class:`opennmt.utils.checkpoint.Checkpoint` instance. If
             not set, no checkpoints will be saved.
         """
-        super().__init__(
-            model, optimizer, checkpoint=checkpoint, is_master=hvd.rank() == 0
-        )
+        super().__init__(model, optimizer, checkpoint=checkpoint)
         self._hvd = hvd
 
+    @property
+    def is_master(self):
+        return self._hvd.rank() == 0
+
     @property
     def num_replicas(self):
         return self._hvd.size()


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
index 88e586d4..adc2d1fc 100644
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -564,6 +564,68 @@ class InputterTest(tf.test.TestCase):
         )
         self.assertIsInstance(dataset, tf.data.Dataset)
 
+    def testBatchAutotuneDataset(self):
+        vocab_file = self._makeTextFile("vocab.txt", ["1", "2", "3", "4"])
+        data_file = self._makeTextFile("data.txt", ["hello world !"])
+        source_inputter = text_inputter.WordEmbedder(embedding_size=10)
+        target_inputter = text_inputter.WordEmbedder(embedding_size=10)
+        target_inputter.set_decoder_mode(mark_start=True, mark_end=True)
+        example_inputter = inputter.ExampleInputter(source_inputter, target_inputter)
+        example_inputter.initialize(
+            {"source_vocabulary": vocab_file, "target_vocabulary": vocab_file}
+        )
+
+        dataset = example_inputter.make_training_dataset(
+            data_file,
+            data_file,
+            batch_size=1024,
+            batch_type="tokens",
+            maximum_features_length=100,
+            maximum_labels_length=120,
+            batch_autotune_mode=True,
+        )
+
+        source, target = next(iter(dataset))
+        self.assertListEqual(source["ids"].shape.as_list(), [8, 100])
+        self.assertListEqual(target["ids"].shape.as_list(), [8, 120])
+        self.assertListEqual(target["ids_out"].shape.as_list(), [8, 120])
+
+    def testBatchAutotuneDatasetMultiSource(self):
+        vocab_file = self._makeTextFile("vocab.txt", ["1", "2", "3", "4"])
+        data_file = self._makeTextFile("data.txt", ["hello world !"])
+        source_inputter = inputter.ParallelInputter(
+            [
+                text_inputter.WordEmbedder(embedding_size=10),
+                text_inputter.WordEmbedder(embedding_size=10),
+            ]
+        )
+        target_inputter = text_inputter.WordEmbedder(embedding_size=10)
+        target_inputter.set_decoder_mode(mark_start=True, mark_end=True)
+        example_inputter = inputter.ExampleInputter(source_inputter, target_inputter)
+        example_inputter.initialize(
+            {
+                "source_1_vocabulary": vocab_file,
+                "source_2_vocabulary": vocab_file,
+                "target_vocabulary": vocab_file,
+            }
+        )
+
+        dataset = example_inputter.make_training_dataset(
+            [data_file, data_file],
+            data_file,
+            batch_size=1024,
+            batch_type="tokens",
+            maximum_features_length=[100, 110],
+            maximum_labels_length=120,
+            batch_autotune_mode=True,
+        )
+
+        source, target = next(iter(dataset))
+        self.assertListEqual(source["inputter_0_ids"].shape.as_list(), [8, 100])
+        self.assertListEqual(source["inputter_1_ids"].shape.as_list(), [8, 110])
+        self.assertListEqual(target["ids"].shape.as_list(), [8, 120])
+        self.assertListEqual(target["ids_out"].shape.as_list(), [8, 120])
+
     def testExampleInputterAsset(self):
         vocab_file = self._makeTextFile("vocab.txt", ["the", "world", "hello", "toto"])
         source_inputter = text_inputter.WordEmbedder(embedding_size=10)