Repository: OpenNMT/OpenNMT-tf

Issue:
Lower throughput than expected in multi-GPU training with online tokenization
There are some performance issues when training on multiple GPUs and enabling the online OpenNMT tokenization. The throughput is lower than expected and GPU usage frequently goes down to 0%. Possible workarounds:

* Tokenize the data before the training
* Use the `SentencePieceTokenizer` that is implemented as a TensorFlow op
* Use Horovod for multi-GPU training

See https://forum.opennmt.net/t/cannot-scale-well-with-multiple-gpus/4239 for a discussion and possible explanation.


Discussion:

Resolving patch:
diff --git a/opennmt/inputters/inputter.py b/opennmt/inputters/inputter.py
--- a/opennmt/inputters/inputter.py
+++ b/opennmt/inputters/inputter.py
@@ -121,13 +121,9 @@ def make_inference_dataset(
         See Also:
           :func:`opennmt.data.inference_pipeline`
         """
-        map_fn = lambda *arg: self.make_features(
-            element=misc.item_or_tuple(arg), training=False
+        transform_fns = _get_dataset_transforms(
+            self, num_threads=num_threads, training=False
         )
-        transform_fns = [
-            lambda dataset: dataset.map(map_fn, num_parallel_calls=num_threads or 1)
-        ]
-
         dataset = self.make_dataset(features_file, training=False)
         dataset = dataset.apply(
             dataset_util.inference_pipeline(
@@ -181,6 +177,30 @@ def get_padded_shapes(self, element_spec, maximum_length=None):
             element_spec,
         )
 
+    def has_prepare_step(self):
+        """Returns ``True`` if this inputter implements a data preparation step
+        in method :meth:`opennmt.inputters.Inputter.prepare_elements`.
+        """
+        return False
+
+    def prepare_elements(self, elements, training=None):
+        """Prepares dataset elements.
+
+        This method is called on a batch of dataset elements. For example, it
+        can be overriden to apply an external pre-tokenization.
+
+        Note that the results of the method are unbatched and then passed to
+        method :meth:`opennmt.inputters.Inputter.make_features`.
+
+        Args:
+          elements: A batch of dataset elements.
+          training: Run in training mode.
+
+        Returns:
+          A (possibly nested) structure of ``tf.Tensor``.
+        """
+        return elements
+
     @abc.abstractmethod
     def make_features(self, element=None, features=None, training=None):
         """Creates features from data.
@@ -305,9 +325,14 @@ def export_assets(self, asset_dir):
             assets.update(inputter.export_assets(asset_dir))
         return assets
 
-    @abc.abstractmethod
-    def make_dataset(self, data_file, training=None):
-        raise NotImplementedError()
+    def has_prepare_step(self):
+        return any(inputter.has_prepare_step() for inputter in self.inputters)
+
+    def prepare_elements(self, elements, training=None):
+        return tuple(
+            inputter.prepare_elements(elts)
+            for inputter, elts in zip(self.inputters, elements)
+        )
 
     def visualize(self, model_root, log_dir):
         for inputter in self.inputters:
@@ -651,13 +676,9 @@ def make_evaluation_dataset(
             data_files = features_file
             length_fn = self.get_length
 
-        map_fn = lambda *arg: self.make_features(
-            element=misc.item_or_tuple(arg), training=False
+        transform_fns = _get_dataset_transforms(
+            self, num_threads=num_threads, training=False
         )
-        transform_fns = [
-            lambda dataset: dataset.map(map_fn, num_parallel_calls=num_threads or 1)
-        ]
-
         dataset = self.make_dataset(data_files, training=False)
         dataset = dataset.apply(
             dataset_util.inference_pipeline(
@@ -753,18 +774,16 @@ def make_training_dataset(
 
         dataset = self.make_dataset(data_files, training=True)
 
-        map_fn = lambda *arg: self.make_features(
-            element=misc.item_or_tuple(arg), training=True
-        )
         filter_fn = lambda *arg: (
             self.keep_for_training(
                 misc.item_or_tuple(arg), maximum_length=maximum_length
             )
         )
-        transform_fns = [
-            lambda dataset: dataset.map(map_fn, num_parallel_calls=num_threads or 4),
-            lambda dataset: dataset.filter(filter_fn),
-        ]
+
+        transform_fns = _get_dataset_transforms(
+            self, num_threads=num_threads, training=True
+        )
+        transform_fns.append(lambda dataset: dataset.filter(filter_fn))
 
         if batch_autotune_mode:
             # In this mode we want to return batches where all sequences are padded
@@ -952,3 +971,32 @@ def make_inference_dataset(
             num_threads=num_threads,
             prefetch_buffer_size=prefetch_buffer_size,
         )
+
+
+def _get_dataset_transforms(
+    inputter,
+    num_threads=None,
+    training=None,
+    prepare_batch_size=128,
+):
+    transform_fns = []
+
+    if inputter.has_prepare_step():
+        prepare_fn = lambda *arg: inputter.prepare_elements(
+            misc.item_or_tuple(arg), training=training
+        )
+        transform_fns.extend(
+            [
+                lambda dataset: dataset.batch(prepare_batch_size),
+                lambda dataset: dataset.map(prepare_fn, num_parallel_calls=num_threads),
+                lambda dataset: dataset.unbatch(),
+            ]
+        )
+
+    map_fn = lambda *arg: inputter.make_features(
+        element=misc.item_or_tuple(arg), training=training
+    )
+    transform_fns.append(
+        lambda dataset: dataset.map(map_fn, num_parallel_calls=num_threads)
+    )
+    return transform_fns
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -269,6 +269,16 @@ def get_dataset_size(self, data_file):
             return list(map(misc.count_lines, data_file))
         return misc.count_lines(data_file)
 
+    def has_prepare_step(self):
+        # For performance reasons, we apply external tokenizers on a batch of
+        # dataset elements during the preparation step.
+        return not self.tokenizer.in_graph and not isinstance(
+            self.tokenizer, tokenizers.SpaceTokenizer
+        )
+
+    def prepare_elements(self, elements, training=None):
+        return {"tokens": self.tokenizer.tokenize(elements, training=training)}
+
     def make_features(self, element=None, features=None, training=None):
         """Tokenizes raw text."""
         self._assert_is_initialized()
@@ -276,10 +286,14 @@ def make_features(self, element=None, features=None, training=None):
             features = {}
         if "tokens" in features:
             return features
-        if "text" in features:
-            element = features.pop("text")
-        element = tf.convert_to_tensor(element, dtype=tf.string)
-        tokens = self.tokenizer.tokenize(element, training=training)
+
+        element = features.pop("text", element)
+        if isinstance(element, dict):
+            tokens = element["tokens"]
+        else:
+            element = tf.convert_to_tensor(element, dtype=tf.string)
+            tokens = self.tokenizer.tokenize(element, training=training)
+
         if isinstance(tokens, tf.RaggedTensor):
             length = tokens.row_lengths()
             tokens = tokens.to_tensor(default_value=constants.PADDING_TOKEN)


Test patch:
diff --git a/opennmt/tests/inputter_test.py b/opennmt/tests/inputter_test.py
--- a/opennmt/tests/inputter_test.py
+++ b/opennmt/tests/inputter_test.py
@@ -654,6 +654,7 @@ def testExampleInputterAsset(self):
             }
         )
         self.assertIsInstance(source_inputter.tokenizer, tokenizers.OpenNMTTokenizer)
+        self.assertTrue(example_inputter.has_prepare_step())
         asset_dir = self.get_temp_dir()
         example_inputter.export_assets(asset_dir)
         self.assertIn("source_tokenizer_config.yml", set(os.listdir(asset_dir)))