save benchmark for 2.11

examples/benchmarking/README.md

@@ -0,0 +1,10 @@
+# 🤗 Benchmark results
+
+Here, you can find a list of the different benchmark results created by the community.
+
+If you would like to list benchmark results on your favorite models of the [model hub](https://huggingface.co/models) here, please open a Pull Request and add it below.
+
+| Benchmark description | Results | Environment info |      Author      |
+|:----------|:-------------|:-------------|------:|
+| PyTorch Benchmark on inference for `bert-base-cased` |[memory](https://github.com/patrickvonplaten/files_to_link_to/blob/master/bert_benchmark/inference_memory.csv) | [env](https://github.com/patrickvonplaten/files_to_link_to/blob/master/bert_benchmark/env.csv) | [Partick von Platen](https://github.com/patrickvonplaten) | 
+| PyTorch Benchmark on inference for `bert-base-cased` |[time](https://github.com/patrickvonplaten/files_to_link_to/blob/master/bert_benchmark/inference_time.csv) | [env](https://github.com/patrickvonplaten/files_to_link_to/blob/master/bert_benchmark/env.csv) | [Partick von Platen](https://github.com/patrickvonplaten) | 

examples/benchmarking/plot_csv_file.py

@@ -1,14 +1,19 @@
 import csv
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Optional
-
-import numpy as np
+from typing import List, Optional
 
 import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.ticker import ScalarFormatter
+
 from transformers import HfArgumentParser
 
 
+def list_field(default=None, metadata=None):
+    return field(default_factory=lambda: default, metadata=metadata)
+
+
 @dataclass
 class PlotArguments:
     """
@@ -24,6 +29,9 @@ class PlotArguments:
         default=False,
         metadata={"help": "Whether the csv file has time results or memory results. Defaults to memory results."},
     )
+    no_log_scale: bool = field(
+        default=False, metadata={"help": "Disable logarithmic scale when plotting"},
+    )
     is_train: bool = field(
         default=False,
         metadata={
@@ -33,6 +41,25 @@ class PlotArguments:
     figure_png_file: Optional[str] = field(
         default=None, metadata={"help": "Filename under which the plot will be saved. If unused no plot is saved."},
     )
+    short_model_names: Optional[List[str]] = list_field(
+        default=None, metadata={"help": "List of model names that are used instead of the ones in the csv file."}
+    )
+
+
+def can_convert_to_int(string):
+    try:
+        int(string)
+        return True
+    except ValueError:
+        return False
+
+
+def can_convert_to_float(string):
+    try:
+        float(string)
+        return True
+    except ValueError:
+        return False
 
 
 class Plot:
@@ -46,16 +73,31 @@ class Plot:
                 model_name = row["model"]
                 self.result_dict[model_name]["bsz"].append(int(row["batch_size"]))
                 self.result_dict[model_name]["seq_len"].append(int(row["sequence_length"]))
-                self.result_dict[model_name]["result"][(int(row["batch_size"]), int(row["sequence_length"]))] = row[
-                    "result"
-                ]
+                if can_convert_to_int(row["result"]):
+                    # value is not None
+                    self.result_dict[model_name]["result"][
+                        (int(row["batch_size"]), int(row["sequence_length"]))
+                    ] = int(row["result"])
+                elif can_convert_to_float(row["result"]):
+                    # value is not None
+                    self.result_dict[model_name]["result"][
+                        (int(row["batch_size"]), int(row["sequence_length"]))
+                    ] = float(row["result"])
 
     def plot(self):
         fig, ax = plt.subplots()
         title_str = "Time usage" if self.args.is_time else "Memory usage"
         title_str = title_str + " for training" if self.args.is_train else title_str + " for inference"
 
-        for model_name in self.result_dict.keys():
+        if not self.args.no_log_scale:
+            # set logarithm scales
+            ax.set_xscale("log")
+            ax.set_yscale("log")
+
+        for axis in [ax.xaxis, ax.yaxis]:
+            axis.set_major_formatter(ScalarFormatter())
+
+        for model_name_idx, model_name in enumerate(self.result_dict.keys()):
             batch_sizes = sorted(list(set(self.result_dict[model_name]["bsz"])))
             sequence_lengths = sorted(list(set(self.result_dict[model_name]["seq_len"])))
             results = self.result_dict[model_name]["result"]
@@ -64,28 +106,33 @@ class Plot:
                 (batch_sizes, sequence_lengths) if self.args.plot_along_batch else (sequence_lengths, batch_sizes)
             )
 
-            plt.xlim(min(x_axis_array), max(x_axis_array))
+            label_model_name = (
+                model_name if self.args.short_model_names is None else self.args.short_model_names[model_name_idx]
+            )
 
             for inner_loop_value in inner_loop_array:
                 if self.args.plot_along_batch:
-                    y_axis_array = np.asarray([results[(x, inner_loop_value)] for x in x_axis_array], dtype=np.int)
+                    y_axis_array = np.asarray(
+                        [results[(x, inner_loop_value)] for x in x_axis_array if (x, inner_loop_value) in results],
+                        dtype=np.int,
+                    )
                 else:
-                    y_axis_array = np.asarray([results[(inner_loop_value, x)] for x in x_axis_array], dtype=np.float32)
-
-                ax.set_xscale("log", basex=2)
-                ax.set_yscale("log", basey=10)
+                    y_axis_array = np.asarray(
+                        [results[(inner_loop_value, x)] for x in x_axis_array if (inner_loop_value, x) in results],
+                        dtype=np.float32,
+                    )
 
                 (x_axis_label, inner_loop_label) = (
-                    ("batch_size", "sequence_length in #tokens")
-                    if self.args.plot_along_batch
-                    else ("sequence_length in #tokens", "batch_size")
+                    ("batch_size", "len") if self.args.plot_along_batch else ("in #tokens", "bsz")
                 )
 
-                x_axis_array = np.asarray(x_axis_array, np.int)
-                plt.scatter(x_axis_array, y_axis_array, label=f"{model_name} - {inner_loop_label}: {inner_loop_value}")
+                x_axis_array = np.asarray(x_axis_array, np.int)[: len(y_axis_array)]
+                plt.scatter(
+                    x_axis_array, y_axis_array, label=f"{label_model_name} - {inner_loop_label}: {inner_loop_value}"
+                )
                 plt.plot(x_axis_array, y_axis_array, "--")
 
-            title_str += f" {model_name} vs."
+            title_str += f" {label_model_name} vs."
 
         title_str = title_str[:-4]
         y_axis_label = "Time in s" if self.args.is_time else "Memory in MB"

examples/benchmarking/run_benchmark.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2018 The HuggingFace Inc. team.
+# Copyright 2020 The HuggingFace Inc. team.
 # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

examples/benchmarking/run_benchmark_tf.py

@@ -0,0 +1,29 @@
+# coding=utf-8
+# Copyright 2018 The HuggingFace Inc. team.
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Benchmarking the library on inference and training in TensorFlow"""
+
+from transformers import HfArgumentParser, TensorFlowBenchmark, TensorFlowBenchmarkArguments
+
+
+def main():
+    parser = HfArgumentParser(TensorFlowBenchmarkArguments)
+    benchmark_args = parser.parse_args_into_dataclasses()[0]
+    benchmark = TensorFlowBenchmark(args=benchmark_args)
+    benchmark.run()
+
+
+if __name__ == "__main__":
+    main()

src/transformers/__init__.py

@@ -78,6 +78,9 @@ from .file_utils import (
     cached_path,
     is_tf_available,
     is_torch_available,
+    is_py3nvml_available,
+    is_torch_tpu_available,
+    is_psutil_available,
 )
 from .hf_argparser import HfArgumentParser
 
@@ -347,7 +350,8 @@ if is_torch_available():
     from .data.datasets import GlueDataset, TextDataset, LineByLineTextDataset, GlueDataTrainingArguments
 
     # Benchmarks
-    from .benchmark import PyTorchBenchmark, PyTorchBenchmarkArguments
+    from .benchmark.benchmark import PyTorchBenchmark
+    from .benchmark.benchmark_args import PyTorchBenchmarkArguments
 
 # TensorFlow
 if is_tf_available():

src/transformers/benchmark/__init__.py

@@ -1,10 +0,0 @@
-# flake8: noqa
-# There's no way to ignore "F401 '...' imported but unused" warnings in this
-# module, but to preserve other warnings. So, don't check this module at all.
-
-from ..file_utils import is_torch_available
-
-
-if is_torch_available():
-    from .benchmark_args import PyTorchBenchmarkArguments
-    from .benchmark import PyTorchBenchmark

src/transformers/benchmark/benchmark.py

@@ -18,13 +18,26 @@
 """
 
 
-import inspect
 import logging
 import timeit
+from typing import Callable, Optional
 
-from transformers import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING, PretrainedConfig, is_torch_available
+from transformers import (
+    MODEL_MAPPING,
+    MODEL_WITH_LM_HEAD_MAPPING,
+    PretrainedConfig,
+    is_py3nvml_available,
+    is_torch_available,
+)
 
-from .benchmark_utils import Benchmark, Memory, start_memory_tracing, stop_memory_tracing
+from .benchmark_utils import (
+    Benchmark,
+    Memory,
+    MemorySummary,
+    measure_peak_memory_cpu,
+    start_memory_tracing,
+    stop_memory_tracing,
+)
 
 
 if is_torch_available():
@@ -32,6 +45,10 @@ if is_torch_available():
     from .benchmark_args import PyTorchBenchmarkArguments
 
 
+if is_py3nvml_available():
+    import py3nvml.py3nvml as nvml
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -45,102 +62,196 @@ class PyTorchBenchmark(Benchmark):
     def framework_version(self):
         return torch.__version__
 
-    def train(self, model_name, batch_size, sequence_length, trace_memory=False):
-        try:
-            config = self.config_dict[model_name]
-            model = MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)
-            model.to(self.args.device)
-            model.train()
+    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
+        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
+        return self._measure_speed(_inference)
 
-            input_ids = torch.randint(
-                model.config.vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
-            )
+    def _inference_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
+        return self._measure_memory(_inference)
 
-            def compute_loss_and_backprob():
-                # TODO: Not all models call labels argument labels => this hack using the function signature should be corrected once all models have a common name for labels
-                function_argument_names = inspect.getfullargspec(model.forward).args
-                if "labels" in function_argument_names:
-                    loss = model(input_ids, labels=input_ids)[0]
-                elif "lm_labels" in function_argument_names:
-                    loss = model(input_ids, lm_labels=input_ids)[0]
-                elif "masked_lm_labels" in function_argument_names:
-                    loss = model(input_ids, masked_lm_labels=input_ids)[0]
-                else:
-                    NotImplementedError(f"{model_name} does not seem to allow training with labels")
+    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
+        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
+        return self._measure_speed(_train)
 
-                loss.backward()
-                model.zero_grad()
+    def _train_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
+        return self._measure_memory(_train)
 
-            if trace_memory is True:
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    trace = start_memory_tracing("transformers")
-                else:
-                    # clear cuda cache
-                    torch.cuda.empty_cache()
-                    torch.cuda.reset_peak_memory_stats()
+    def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
+        config = self.config_dict[model_name]
 
-                # calculate loss and do backpropagation
-                compute_loss_and_backprob()
+        if self.args.torchscript:
+            config.torchscript = True
 
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    summary = stop_memory_tracing(trace)
-                    memory = summary.total
-                else:
-                    memory = Memory(torch.cuda.max_memory_reserved())
-
-                return memory
-            else:
-                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
-                runtimes = timeit.repeat(lambda: compute_loss_and_backprob(), repeat=self.args.repeat, number=10,)
-                return min(runtimes) / 10.0
-        except RuntimeError as e:
-            self.print_fn("Doesn't fit on GPU. {}".format(e))
-            return "N/A"
-
-    def inference(self, model_name, batch_size, sequence_length, trace_memory=False):
-        try:
-            config = self.config_dict[model_name]
+        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        if not self.args.only_pretrain_model and has_model_class_in_config:
+            try:
+                model_class = config.architectures[0]
+                transformers_module = __import__("transformers", fromlist=[model_class])
+                model_cls = getattr(transformers_module, model_class)
+                model = model_cls(config)
+            except ImportError:
+                raise ImportError(
+                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`."
+                )
+        else:
             model = MODEL_MAPPING[config.__class__](config)
-            model.to(self.args.device)
-            model.eval()
 
-            input_ids = torch.randint(
-                config.vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
-            )
-            if trace_memory is True:
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    trace = start_memory_tracing("transformers")
-                else:
-                    # clear cuda cache
-                    torch.cuda.empty_cache()
-                    if hasattr(torch.cuda, "max_memory_reserved"):
-                        torch.cuda.reset_peak_memory_stats()
-                    else:
-                        logger.info(
-                            "Please consider updating PyTorch to version 1.4 to get more accuracy on GPU memory usage"
-                        )
-                        torch.cuda.reset_max_memory_cached()
+        model.eval()
+        model.to(self.args.device)
 
-                model(input_ids)
+        # encoder-decoder has vocab size saved differently
+        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
+        input_ids = torch.randint(vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device)
 
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    summary = stop_memory_tracing(trace)
-                    memory = summary.total
-                else:
-                    if hasattr(torch.cuda, "max_memory_reserved"):
-                        memory = Memory(torch.cuda.max_memory_reserved())
-                    else:
-                        logger.info(
-                            "Please consider updating PyTorch to version 1.4 to get more accuracy on GPU memory usage"
-                        )
-                        memory = Memory(torch.cuda.max_memory_cached())
+        if self.args.fp16:
+            logger.info("Running training in Mixed Precision...")
+            assert self.args.is_gpu, "Mixed precision is possible only for GPU."
+            # amp seems to have memory leaks so that memory usage
+            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
+            model.half()
 
-                return memory
-            else:
-                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
-                runtimes = timeit.repeat(lambda: model(input_ids), repeat=self.args.repeat, number=10,)
-                return min(runtimes) / 10.0
+        if self.args.torchscript:
+            with torch.no_grad():
+                inference_model = torch.jit.trace(model, input_ids)
+        else:
+            inference_model = model
 
+        def encoder_decoder_forward():
+            with torch.no_grad():
+                outputs = inference_model(input_ids, decoder_input_ids=input_ids)
+            return outputs
+
+        def encoder_forward():
+            with torch.no_grad():
+                outputs = inference_model(input_ids)
+            return outputs
+
+        _forward = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
+        return _forward
+
+    def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
+        config = self.config_dict[model_name]
+
+        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        if not self.args.only_pretrain_model and has_model_class_in_config:
+            try:
+                model_class = config.architectures[0]
+                transformers_module = __import__("transformers", fromlist=[model_class])
+                model_cls = getattr(transformers_module, model_class)
+                model = model_cls(config)
+            except ImportError:
+                raise ImportError(
+                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`."
+                )
+        else:
+            model = MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)
+
+        if self.args.torchscript:
+            raise NotImplementedError("Training for torchscript is currently not implemented")
+        else:
+            train_model = model
+
+        model.train()
+        model.to(self.args.device)
+
+        # encoder-decoder has vocab size saved differently
+        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
+        input_ids = torch.randint(vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device)
+
+        if self.args.fp16:
+            logger.info("Running training in Mixed Precision...")
+            assert self.args.is_gpu, "Mixed precision is possible only for GPU."
+
+            # amp seems to have memory leaks so that memory usage
+            # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
+            model.half()
+
+        def compute_loss_and_backprob_encoder():
+            loss = train_model(input_ids, labels=input_ids)[0]
+            loss.backward()
+            return loss
+
+        def compute_loss_and_backprob_encoder_decoder():
+            loss = train_model(input_ids, decoder_input_ids=input_ids, labels=input_ids)[0]
+            loss.backward()
+            return loss
+
+        _train = (
+            compute_loss_and_backprob_encoder_decoder
+            if config.is_encoder_decoder
+            else compute_loss_and_backprob_encoder
+        )
+        return _train
+
+    def _measure_speed(self, func) -> float:
+        try:
+            if self.args.is_tpu or self.args.torchscript:
+                # run additional 10 times to stabilize compilation for tpu and torchscript
+                logger.info("Do inference on TPU or torchscript. Running model 5 times to stabilize compilation")
+                timeit.repeat(
+                    func, repeat=1, number=5,
+                )
+
+            # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
+            runtimes = timeit.repeat(func, repeat=self.args.repeat, number=10,)
+
+            if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
+                import torch_xla.debug.metrics as met
+
+                self.print_fn(met.metrics_report())
+
+            return min(runtimes) / 10.0
         except RuntimeError as e:
             self.print_fn("Doesn't fit on GPU. {}".format(e))
             return "N/A"
+
+    def _measure_memory(self, func: Callable[[], None]) -> [Memory, MemorySummary]:
+        try:
+            if self.args.trace_memory_line_by_line:
+                trace = start_memory_tracing("transformers")
+
+            if self.args.is_tpu:
+                # tpu
+                raise NotImplementedError(
+                    "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with `--no_memory` or `args.no_memory=True`"
+                )
+            elif self.args.is_gpu:
+                if not is_py3nvml_available():
+                    logger.warning(
+                        "py3nvml not installed, we won't log GPU memory usage. "
+                        "Install py3nvml (pip install py3nvml) to log information about GPU."
+                    )
+                    memory = "N/A"
+                else:
+                    logger.info(
+                        "Measuring total GPU usage on GPU device. Make sure to not have additional processes running on the same GPU."
+                    )
+                    # init nvml
+                    nvml.nvmlInit()
+                    func()
+                    handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
+                    meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
+                    max_bytes_in_use = meminfo.used
+                    memory = Memory(max_bytes_in_use)
+                    # shutdown nvml
+                    nvml.nvmlShutdown()
+            else:
+                # cpu
+                memory_bytes = measure_peak_memory_cpu(func)
+                memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
+
+            if self.args.trace_memory_line_by_line:
+                summary = stop_memory_tracing(trace)
+            else:
+                summary = None
+
+            return memory, summary
+        except RuntimeError as e:
+            self.print_fn("Doesn't fit on GPU. {}".format(e))
+            return "N/A", None

src/transformers/benchmark/benchmark_args.py

@@ -18,34 +18,33 @@ import logging
 from dataclasses import dataclass, field
 from typing import Tuple
 
-from ..file_utils import cached_property, is_torch_available, torch_required
+from ..file_utils import cached_property, is_torch_available, is_torch_tpu_available, torch_required
 from .benchmark_args_utils import BenchmarkArguments
 
 
 if is_torch_available():
     import torch
 
-try:
+if is_torch_tpu_available():
     import torch_xla.core.xla_model as xm
 
-    _has_tpu = True
-except ImportError:
-    _has_tpu = False
-
-
-@torch_required
-def is_tpu_available():
-    return _has_tpu
-
 
 logger = logging.getLogger(__name__)
 
 
 @dataclass
 class PyTorchBenchmarkArguments(BenchmarkArguments):
-    no_cuda: bool = field(default=False, metadata={"help": "Whether to run on available cuda devices"})
     torchscript: bool = field(default=False, metadata={"help": "Trace the models using torchscript"})
-    fp16: bool = field(default=False, metadata={"help": "Use FP16 to accelerate inference."})
+    torch_xla_tpu_print_metrics: bool = field(default=False, metadata={"help": "Print Xla/PyTorch tpu metrics"})
+    fp16_opt_level: str = field(
+        default="O1",
+        metadata={
+            "help": (
+                "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
+                "See details at https://nvidia.github.io/apex/amp.html"
+            )
+        },
+    )
 
     @cached_property
     @torch_required
@@ -54,7 +53,7 @@ class PyTorchBenchmarkArguments(BenchmarkArguments):
         if self.no_cuda:
             device = torch.device("cpu")
             n_gpu = 0
-        elif is_tpu_available():
+        elif is_torch_tpu_available():
             device = xm.xla_device()
             n_gpu = 0
         else:
@@ -62,9 +61,14 @@ class PyTorchBenchmarkArguments(BenchmarkArguments):
             n_gpu = torch.cuda.device_count()
         return device, n_gpu
 
+    @property
+    def is_tpu(self):
+        return is_torch_tpu_available() and not self.no_tpu
+
     @property
     @torch_required
     def device_idx(self) -> int:
+        # TODO(PVP): currently only single GPU is supported
         return torch.cuda.current_device()
 
     @property
@@ -76,3 +80,7 @@ class PyTorchBenchmarkArguments(BenchmarkArguments):
     @torch_required
     def n_gpu(self):
         return self._setup_devices[1]
+
+    @property
+    def is_gpu(self):
+        return self.n_gpu > 0

src/transformers/benchmark/benchmark_args_tf.py

@@ -0,0 +1,105 @@
+# coding=utf-8
+# Copyright 2018 The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from ..file_utils import cached_property, is_tf_available, tf_required
+from .benchmark_args_utils import BenchmarkArguments
+
+
+if is_tf_available():
+    import tensorflow as tf
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TensorFlowBenchmarkArguments(BenchmarkArguments):
+    tpu_name: str = field(
+        default=None, metadata={"help": "Name of TPU"},
+    )
+    device_idx: int = field(
+        default=0, metadata={"help": "CPU / GPU device index. Defaults to 0."},
+    )
+    eager_mode: bool = field(default=False, metadata={"help": "Benchmark models in eager model."})
+    use_xla: bool = field(
+        default=False,
+        metadata={
+            "help": "Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`."
+        },
+    )
+
+    @cached_property
+    @tf_required
+    def _setup_tpu(self) -> Tuple["tf.distribute.cluster_resolver.TPUClusterResolver"]:
+        if not self.no_tpu:
+            try:
+                if self.tpu_name:
+                    tpu = tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name)
+                else:
+                    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
+            except ValueError:
+                tpu = None
+        return tpu
+
+    @cached_property
+    @tf_required
+    def _setup_strategy(self) -> Tuple["tf.distribute.Strategy", "tf.distribute.cluster_resolver.TPUClusterResolver"]:
+        if self.is_tpu:
+            tf.config.experimental_connect_to_cluster(self._setup_tpu)
+            tf.tpu.experimental.initialize_tpu_system(self._setup_tpu)
+
+            strategy = tf.distribute.experimental.TPUStrategy(self._setup_tpu)
+        else:
+            # currently no multi gpu is allowed
+            if self.is_gpu:
+                # TODO: Currently only single GPU is supported
+                tf.config.experimental.set_visible_devices(self.gpu_list[self.device_idx], "GPU")
+                strategy = tf.distribute.OneDeviceStrategy(device=f"/gpu:{self.device_idx}")
+            else:
+                tf.config.experimental.set_visible_devices([], "GPU")  # disable GPU
+                strategy = tf.distribute.OneDeviceStrategy(device=f"/cpu:{self.device_idx}")
+
+        return strategy
+
+    @property
+    @tf_required
+    def is_tpu(self) -> bool:
+        return self._setup_tpu is not None
+
+    @property
+    @tf_required
+    def strategy(self) -> "tf.distribute.Strategy":
+        return self._setup_strategy
+
+    @property
+    @tf_required
+    def gpu_list(self):
+        return tf.config.list_physical_devices("GPU")
+
+    @property
+    @tf_required
+    def n_gpu(self) -> int:
+        if not self.no_cuda:
+            return len(self.gpu_list)
+        return 0
+
+    @property
+    def is_gpu(self) -> bool:
+        return self.n_gpu > 0

src/transformers/benchmark/benchmark_args_utils.py

@@ -16,11 +16,15 @@
 
 import dataclasses
 import json
+import logging
 from dataclasses import dataclass, field
 from time import time
 from typing import List
 
 
+logger = logging.getLogger(__name__)
+
+
 def list_field(default=None, metadata=None):
     return field(default_factory=lambda: default, metadata=metadata)
 
@@ -53,6 +57,9 @@ class BenchmarkArguments:
     )
 
     no_inference: bool = field(default=False, metadata={"help": "Don't benchmark inference of model"})
+    no_cuda: bool = field(default=False, metadata={"help": "Whether to run on available cuda devices"})
+    no_tpu: bool = field(default=False, metadata={"help": "Whether to run on available tpu devices"})
+    fp16: bool = field(default=False, metadata={"help": "Use FP16 to accelerate inference."})
     training: bool = field(default=False, metadata={"help": "Benchmark training of model"})
     verbose: bool = field(default=False, metadata={"help": "Verbose memory tracing"})
     no_speed: bool = field(default=False, metadata={"help": "Don't perform speed measurments"})
@@ -61,6 +68,12 @@ class BenchmarkArguments:
     save_to_csv: bool = field(default=False, metadata={"help": "Save result to a CSV file"})
     log_print: bool = field(default=False, metadata={"help": "Save all print statements in a log file"})
     no_env_print: bool = field(default=False, metadata={"help": "Don't print environment information"})
+    no_multi_process: bool = field(
+        default=False,
+        metadata={
+            "help": "Don't use multiprocessing for memory and speed measurement. It is highly recommended to use multiprocessing for accurate CPU and GPU memory measurements. This option should only be used for debugging / testing and on TPU."
+        },
+    )
     inference_time_csv_file: str = field(
         default=f"inference_time_{round(time())}.csv",
         metadata={"help": "CSV filename used if saving time results to csv."},
@@ -86,6 +99,12 @@ class BenchmarkArguments:
         metadata={"help": "Log filename used if print statements are saved in log."},
     )
     repeat: int = field(default=3, metadata={"help": "Times an experiment will be run."})
+    only_pretrain_model: bool = field(
+        default=False,
+        metadata={
+            "help": "Instead of loading the model as defined in `config.architectures` if exists, just load the pretrain model weights."
+        },
+    )
 
     def to_json_string(self):
         """
@@ -95,4 +114,17 @@ class BenchmarkArguments:
 
     @property
     def model_names(self):
+        assert (
+            len(self.models) > 0
+        ), "Please make sure you provide at least one model name / model identifier, *e.g.* `--models bert-base-cased` or `args.models = ['bert-base-cased']."
         return self.models
+
+    @property
+    def do_multi_processing(self):
+        if self.no_multi_process:
+            return False
+        elif self.is_tpu:
+            logger.info("Multiprocessing is currently not possible on TPU.")
+            return False
+        else:
+            return True

src/transformers/benchmark/benchmark_tf.py

@@ -0,0 +1,286 @@
+# coding=utf-8
+# Copyright 2018 The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+    Benchmarking the library on inference and training in PyTorch.
+"""
+
+
+import logging
+import random
+import timeit
+from functools import wraps
+from typing import Callable, Optional
+
+from transformers import (
+    TF_MODEL_MAPPING,
+    TF_MODEL_WITH_LM_HEAD_MAPPING,
+    PretrainedConfig,
+    is_py3nvml_available,
+    is_tf_available,
+)
+
+from .benchmark_utils import (
+    Benchmark,
+    Memory,
+    MemorySummary,
+    measure_peak_memory_cpu,
+    start_memory_tracing,
+    stop_memory_tracing,
+)
+
+
+if is_tf_available():
+    import tensorflow as tf
+    from .benchmark_args_tf import TensorFlowBenchmarkArguments
+    from tensorflow.python.framework.errors_impl import ResourceExhaustedError
+
+if is_py3nvml_available():
+    import py3nvml.py3nvml as nvml
+
+logger = logging.getLogger(__name__)
+
+
+def run_with_tf_optimizations(do_eager_mode: bool, use_xla: bool):
+    def run_func(func):
+        @wraps(func)
+        def run_in_eager_mode(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        @wraps(func)
+        @tf.function(experimental_compile=use_xla)
+        def run_in_graph_mode(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        if do_eager_mode is True:
+            assert (
+                use_xla is False
+            ), "Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`."
+            return run_in_eager_mode
+        else:
+            return run_in_graph_mode
+
+    return run_func
+
+
+def random_input_ids(batch_size: int, sequence_length: int, vocab_size: int) -> ["tf.Tensor"]:
+    rng = random.Random()
+    values = [rng.randint(0, vocab_size - 1) for i in range(batch_size * sequence_length)]
+    return tf.constant(values, shape=(batch_size, sequence_length), dtype=tf.int32)
+
+
+class TensorFlowBenchmark(Benchmark):
+
+    args: TensorFlowBenchmarkArguments
+    configs: PretrainedConfig
+    framework: str = "TensorFlow"
+
+    @property
+    def framework_version(self):
+        return tf.__version__
+
+    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
+        # initialize GPU on separate process
+        strategy = self.args.strategy
+        assert strategy is not None, "A device strategy has to be initialized before using TensorFlow."
+        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
+        return self._measure_speed(_inference)
+
+    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
+        strategy = self.args.strategy
+        assert strategy is not None, "A device strategy has to be initialized before using TensorFlow."
+        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
+        return self._measure_speed(_train)
+
+    def _inference_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        # initialize GPU on separate process
+        if self.args.is_gpu:
+            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx], True)
+        strategy = self.args.strategy
+        assert strategy is not None, "A device strategy has to be initialized before using TensorFlow."
+        _inference = self._prepare_inference_func(model_name, batch_size, sequence_length)
+        return self._measure_memory(_inference)
+
+    def _train_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        if self.args.is_gpu:
+            tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx], True)
+        strategy = self.args.strategy
+        assert strategy is not None, "A device strategy has to be initialized before using TensorFlow."
+
+        _train = self._prepare_train_func(model_name, batch_size, sequence_length)
+        return self._measure_memory(_train)
+
+    def _prepare_inference_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
+        config = self.config_dict[model_name]
+
+        if self.args.fp16:
+            raise NotImplementedError("Mixed precision is currently not supported.")
+
+        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        if not self.args.only_pretrain_model and has_model_class_in_config:
+            try:
+                model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
+                transformers_module = __import__("transformers", fromlist=[model_class])
+                model_cls = getattr(transformers_module, model_class)
+                model = model_cls(config)
+            except ImportError:
+                raise ImportError(
+                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`."
+                )
+        else:
+            model = TF_MODEL_MAPPING[config.__class__](config)
+
+        # encoder-decoder has vocab size saved differently
+        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
+        input_ids = random_input_ids(batch_size, sequence_length, vocab_size)
+
+        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
+        def encoder_decoder_forward():
+            return model(input_ids, decoder_input_ids=input_ids, training=False)
+
+        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
+        def encoder_forward():
+            return model(input_ids, training=False)
+
+        _inference = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
+
+        return _inference
+
+    def _prepare_train_func(self, model_name: str, batch_size: int, sequence_length: int) -> Callable[[], None]:
+        config = self.config_dict[model_name]
+
+        assert (
+            self.args.eager_mode is False
+        ), "Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`."
+
+        if self.args.fp16:
+            raise NotImplementedError("Mixed precision is currently not supported.")
+
+        has_model_class_in_config = hasattr(config, "architecture") and len(config.architectures) > 1
+        if not self.args.only_pretrain_model and has_model_class_in_config:
+            try:
+                model_class = "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
+                transformers_module = __import__("transformers", fromlist=[model_class])
+                model_cls = getattr(transformers_module, model_class)
+                model = model_cls(config)
+            except ImportError:
+                raise ImportError(
+                    f"{model_class} does not exist. If you just want to test the pretrained model, you might want to set `--only_pretrain_model` or `args.only_pretrain_model=True`."
+                )
+        else:
+            model = TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)
+
+        # encoder-decoder has vocab size saved differently
+        vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
+        input_ids = random_input_ids(batch_size, sequence_length, vocab_size)
+
+        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
+        def encoder_decoder_train():
+            loss = model(input_ids, decoder_input_ids=input_ids, labels=input_ids, training=True)[0]
+            gradients = tf.gradients(loss, model.trainable_variables)
+            return gradients
+
+        @run_with_tf_optimizations(self.args.eager_mode, self.args.use_xla)
+        def encoder_train():
+            loss = model(input_ids, labels=input_ids, training=True)[0]
+            gradients = tf.gradients(loss, model.trainable_variables)
+            return gradients
+
+        _train = encoder_decoder_train if config.is_encoder_decoder else encoder_train
+
+        return _train
+
+    def _measure_speed(self, func) -> float:
+        with self.args.strategy.scope():
+            try:
+                if self.args.is_tpu or self.args.use_xla:
+                    # run additional 10 times to stabilize compilation for tpu
+                    logger.info("Do inference on TPU. Running model 5 times to stabilize compilation")
+                    timeit.repeat(func, repeat=1, number=5)
+
+                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
+                runtimes = timeit.repeat(func, repeat=self.args.repeat, number=10,)
+
+                return min(runtimes) / 10.0
+            except ResourceExhaustedError as e:
+                self.print_fn("Doesn't fit on GPU. {}".format(e))
+
+    def _measure_memory(self, func: Callable[[], None]) -> [Memory, MemorySummary]:
+        logger.info(
+            "Note that TensorFlow allocates more memory than"
+            "it might need to speed up computation."
+            "The memory reported here corresponds to the memory"
+            "reported by `nvidia-smi`, which can vary depending"
+            "on total available memory on the GPU that is used."
+        )
+        with self.args.strategy.scope():
+            try:
+                if self.args.trace_memory_line_by_line:
+                    assert (
+                        self.args.eager_mode
+                    ), "`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory consumption line by line."
+                    trace = start_memory_tracing("transformers")
+
+                if self.args.is_tpu:
+                    # tpu
+                    raise NotImplementedError(
+                        "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with `args.no_memory=True`"
+                    )
+                elif self.args.is_gpu:
+                    # gpu
+                    if not is_py3nvml_available():
+                        logger.warning(
+                            "py3nvml not installed, we won't log GPU memory usage. "
+                            "Install py3nvml (pip install py3nvml) to log information about GPU."
+                        )
+                        memory = "N/A"
+                    else:
+                        logger.info(
+                            "Measuring total GPU usage on GPU device. Make sure to not have additional processes running on the same GPU."
+                        )
+                        # init nvml
+                        nvml.nvmlInit()
+                        func()
+                        handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
+                        meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
+                        max_bytes_in_use = meminfo.used
+                        memory = Memory(max_bytes_in_use)
+                        # shutdown nvml
+                        nvml.nvmlShutdown()
+                else:
+                    # cpu
+                    if self.args.trace_memory_line_by_line:
+                        logger.info(
+                            "When enabling line by line tracing, the max peak memory for CPU is inaccurate in TensorFlow."
+                        )
+                        memory = None
+                    else:
+                        memory_bytes = measure_peak_memory_cpu(func)
+                        memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
+                if self.args.trace_memory_line_by_line:
+                    summary = stop_memory_tracing(trace)
+                    if memory is None:
+                        memory = summary.total
+                else:
+                    summary = None
+
+                return memory, summary
+            except ResourceExhaustedError as e:
+                self.print_fn("Doesn't fit on GPU. {}".format(e))
+                return "N/A", None

src/transformers/benchmark/benchmark_utils.py

@@ -14,12 +14,14 @@ import sys
 from abc import ABC, abstractmethod
 from collections import defaultdict, namedtuple
 from datetime import datetime
-from typing import Iterable, List, NamedTuple, Optional, Union
+from multiprocessing import Pipe, Process, Queue
+from multiprocessing.connection import Connection
+from typing import Callable, Iterable, List, NamedTuple, Optional, Union
 
 from transformers import AutoConfig, PretrainedConfig
 from transformers import __version__ as version
 
-from ..file_utils import is_tf_available, is_torch_available
+from ..file_utils import is_psutil_available, is_py3nvml_available, is_tf_available, is_torch_available
 from .benchmark_args_utils import BenchmarkArguments
 
 
@@ -29,6 +31,17 @@ if is_torch_available():
 if is_tf_available():
     from tensorflow.python.eager import context as tf_context
 
+if is_psutil_available():
+    import psutil
+
+if is_py3nvml_available():
+    import py3nvml.py3nvml as nvml
+
+if platform.system() == "Windows":
+    from signal import CTRL_C_EVENT as SIGKILL
+else:
+    from signal import SIGKILL
+
 
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 
@@ -36,10 +49,57 @@ logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 _is_memory_tracing_enabled = False
 
 BenchmarkOutput = namedtuple(
-    "BenchmarkOutput", ["time_inference_result", "memory_inference_result", "time_train_result", "memory_train_result"]
+    "BenchmarkOutput",
+    [
+        "time_inference_result",
+        "memory_inference_result",
+        "time_train_result",
+        "memory_train_result",
+        "inference_summary",
+        "train_summary",
+    ],
 )
 
 
+def separate_process_wrapper_fn(func: Callable[[], None], do_multi_processing: bool) -> Callable[[], None]:
+    """
+        This function wraps another function into its own separated process.
+        In order to ensure accurate memory measurements it is important that the function
+        is executed in a separate process
+
+        Args:
+            - `func`: (`callable`): function() -> ...
+                generic function which will be executed in its own separate process
+            - `do_multi_processing`: (`bool`)
+                Whether to run function on separate process or not
+    """
+
+    def multi_process_func(*args, **kwargs):
+        # run function in an individual
+        # process to get correct memory
+        def wrapper_func(queue: Queue, *args):
+            try:
+                result = func(*args)
+            except Exception as e:
+                logger.error(e)
+                print(e)
+                result = "N/A"
+            queue.put(result)
+
+        queue = Queue()
+        p = Process(target=wrapper_func, args=[queue] + list(args))
+        p.start()
+        result = queue.get()
+        p.join()
+        return result
+
+    if do_multi_processing:
+        logging.info("fFunction {func} is executed in its own process...")
+        return multi_process_func
+    else:
+        return func
+
+
 def is_memory_tracing_enabled():
     global _is_memory_tracing_enabled
     return _is_memory_tracing_enabled
@@ -120,6 +180,129 @@ class MemorySummary(NamedTuple):
 MemoryTrace = List[UsedMemoryState]
 
 
+def measure_peak_memory_cpu(function: Callable[[], None], interval=0.5, device_idx=None) -> int:
+    """
+        measures peak cpu memory consumption of a given `function`
+        running the function for at least interval seconds
+        and at most 20 * interval seconds.
+        This function is heavily inspired by: `memory_usage`
+        of the package `memory_profiler`: https://github.com/pythonprofilers/memory_profiler/blob/895c4ac7a08020d66ae001e24067da6dcea42451/memory_profiler.py#L239
+
+        Args:
+            - `function`: (`callable`): function() -> ...
+                function without any arguments to measure for which to measure the peak memory
+
+            - `interval`: (`float`, `optional`, defaults to `0.5`)
+                interval in second for which to measure the memory usage
+
+            - `device_idx`: (`int`, `optional`, defaults to `None`)
+                device id for which to measure gpu usage
+
+        Returns:
+            - `max_memory`: (`int`)
+                cosumed memory peak in Bytes
+    """
+
+    def get_cpu_memory(process_id: int) -> int:
+        """
+            measures current cpu memory usage of a given `process_id`
+
+            Args:
+                - `process_id`: (`int`)
+                    process_id for which to measure memory
+
+            Returns
+                - `memory`: (`int`)
+                    cosumed memory in Bytes
+        """
+        process = psutil.Process(process_id)
+        try:
+            meminfo_attr = "memory_info" if hasattr(process, "memory_info") else "get_memory_info"
+            memory = getattr(process, meminfo_attr)()[0]
+        except psutil.AccessDenied:
+            raise ValueError("Error with Psutil.")
+        return memory
+
+    if not is_psutil_available():
+        logger.warning(
+            "Psutil not installed, we won't log CPU memory usage. "
+            "Install Psutil (pip install psutil) to use CPU memory tracing."
+        )
+        max_memory = "N/A"
+    else:
+
+        class MemoryMeasureProcess(Process):
+
+            """
+                `MemoryMeasureProcess` inherits from `Process` and overwrites
+                its `run()` method. Used to measure the memory usage of a process
+            """
+
+            def __init__(self, process_id: int, child_connection: Connection, interval: float):
+                super().__init__()
+                self.process_id = process_id
+                self.interval = interval
+                self.connection = child_connection
+                self.num_measurements = 1
+                self.mem_usage = get_cpu_memory(self.process_id)
+
+            def run(self):
+                self.connection.send(0)
+                stop = False
+                while True:
+                    self.mem_usage = max(self.mem_usage, get_cpu_memory(self.process_id))
+                    self.num_measurements += 1
+
+                    if stop:
+                        break
+
+                    stop = self.connection.poll(self.interval)
+
+                # send results to parent pipe
+                self.connection.send(self.mem_usage)
+                self.connection.send(self.num_measurements)
+
+        while True:
+            # create child, parent connection
+            child_connection, parent_connection = Pipe()
+
+            # instantiate process
+            mem_process = MemoryMeasureProcess(os.getpid(), child_connection, interval)
+            mem_process.start()
+
+            # wait until we get memory
+            parent_connection.recv()
+
+            try:
+                # execute function
+                function()
+
+                # start parent connection
+                parent_connection.send(0)
+
+                # receive memory and num measurements
+                max_memory = parent_connection.recv()
+                num_measurements = parent_connection.recv()
+            except Exception:
+                # kill process in a clean way
+                parent = psutil.Process(os.getpid())
+                for child in parent.children(recursive=True):
+                    os.kill(child.pid, SIGKILL)
+                mem_process.join(0)
+                raise RuntimeError("Process killed. Error in Process")
+
+            # run process at least 20 * interval or until it finishes
+            mem_process.join(20 * interval)
+
+            if (num_measurements > 4) or (interval < 1e-6):
+                break
+
+            # reduce interval
+            interval /= 10
+
+        return max_memory
+
+
 def start_memory_tracing(
     modules_to_trace: Optional[Union[str, Iterable[str]]] = None,
     modules_not_to_trace: Optional[Union[str, Iterable[str]]] = None,
@@ -159,34 +342,31 @@ def start_memory_tracing(
             - 'line_text' (string): Text of the line in the python script
 
     """
-    try:
-        import psutil
-    except (ImportError):
+    if is_psutil_available():
+        process = psutil.Process(os.getpid())
+    else:
         logger.warning(
             "Psutil not installed, we won't log CPU memory usage. "
             "Install psutil (pip install psutil) to use CPU memory tracing."
         )
         process = None
+
+    if is_py3nvml_available():
+        try:
+            nvml.nvmlInit()
+            devices = list(range(nvml.nvmlDeviceGetCount())) if gpus_to_trace is None else gpus_to_trace
+            nvml.nvmlShutdown()
+        except (OSError, nvml.NVMLError):
+            logger.warning("Error while initializing comunication with GPU. " "We won't perform GPU memory tracing.")
+            log_gpu = False
+        else:
+            log_gpu = is_torch_available() or is_tf_available()
     else:
-        process = psutil.Process(os.getpid())
-
-    try:
-        from py3nvml import py3nvml
-
-        py3nvml.nvmlInit()
-        devices = list(range(py3nvml.nvmlDeviceGetCount())) if gpus_to_trace is None else gpus_to_trace
-        py3nvml.nvmlShutdown()
-    except ImportError:
         logger.warning(
             "py3nvml not installed, we won't log GPU memory usage. "
             "Install py3nvml (pip install py3nvml) to use GPU memory tracing."
         )
         log_gpu = False
-    except (OSError, py3nvml.NVMLError):
-        logger.warning("Error while initializing comunication with GPU. " "We won't perform GPU memory tracing.")
-        log_gpu = False
-    else:
-        log_gpu = is_torch_available() or is_tf_available()
 
     memory_trace = []
 
@@ -206,6 +386,9 @@ def start_memory_tracing(
             elif isinstance(events_to_trace, (list, tuple)) and event not in events_to_trace:
                 return traceit
 
+        if "__name__" not in frame.f_globals:
+            return traceit
+
         # Filter modules
         name = frame.f_globals["__name__"]
         if not isinstance(name, str):
@@ -248,14 +431,14 @@ def start_memory_tracing(
                 tf_context.context()._clear_caches()  # See https://github.com/tensorflow/tensorflow/issues/20218#issuecomment-416771802
 
             # Sum used memory for all GPUs
-            py3nvml.nvmlInit()
+            nvml.nvmlInit()
 
             for i in devices:
-                handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
-                meminfo = py3nvml.nvmlDeviceGetMemoryInfo(handle)
+                handle = nvml.nvmlDeviceGetHandleByIndex(i)
+                meminfo = nvml.nvmlDeviceGetMemoryInfo(handle)
                 gpu_mem += meminfo.used
 
-            py3nvml.nvmlShutdown()
+            nvml.nvmlShutdown()
 
         mem_state = UsedMemoryState(traced_state, cpu_mem, gpu_mem)
         memory_trace.append(mem_state)
@@ -385,7 +568,6 @@ class Benchmark(ABC):
 
     def __init__(self, args: BenchmarkArguments = None, configs: PretrainedConfig = None):
         self.args = args
-
         if configs is None:
             self.config_dict = {
                 model_name: AutoConfig.from_pretrained(model_name) for model_name in self.args.model_names
@@ -393,6 +575,11 @@ class Benchmark(ABC):
         else:
             self.config_dict = {model_name: config for model_name, config in zip(self.args.model_names, configs)}
 
+        if not self.args.no_memory and os.getenv("TRANSFORMERS_USE_MULTIPROCESSING") == 0:
+            logger.warning(
+                "Memory consumption will not be measured accurately if `args.no_multi_process` is set to `True.` The flag 'TRANSFORMERS_USE_MULTIPROCESSING' should only be disabled for debugging / testing."
+            )
+
         self._print_fn = None
         self._framework_version = None
         self._environment_info = None
@@ -401,15 +588,10 @@ class Benchmark(ABC):
     def print_fn(self):
         if self._print_fn is None:
             if self.args.log_print:
-                logging.basicConfig(
-                    level=logging.DEBUG,
-                    filename=self.args.log_filename,
-                    filemode="a+",
-                    format="%(asctime)-15s %(levelname)-8s %(message)s",
-                )
 
                 def print_and_log(*args):
-                    logging.info(*args)
+                    with open(self.args.log_filename, "a") as log_file:
+                        log_file.write("".join(args) + "\n")
                     print(*args)
 
                 self._print_fn = print_and_log
@@ -417,23 +599,43 @@ class Benchmark(ABC):
                 self._print_fn = print
         return self._print_fn
 
-    @property
-    def is_gpu(self):
-        return self.args.n_gpu > 0
-
     @property
     @abstractmethod
     def framework_version(self):
         pass
 
     @abstractmethod
-    def train(self, model_name, batch_size, sequence_length):
+    def _inference_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
         pass
 
     @abstractmethod
-    def inference(self, model_name, batch_size, sequence_length):
+    def _train_speed(self, model_name: str, batch_size: int, sequence_length: int) -> float:
         pass
 
+    @abstractmethod
+    def _inference_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        pass
+
+    @abstractmethod
+    def _train_memory(
+        self, model_name: str, batch_size: int, sequence_length: int
+    ) -> [Memory, Optional[MemorySummary]]:
+        pass
+
+    def inference_speed(self, *args, **kwargs) -> float:
+        return separate_process_wrapper_fn(self._inference_speed, self.args.do_multi_processing)(*args, **kwargs)
+
+    def train_speed(self, *args, **kwargs) -> float:
+        return separate_process_wrapper_fn(self._train_speed, self.args.do_multi_processing)(*args, **kwargs)
+
+    def inference_memory(self, *args, **kwargs) -> [Memory, Optional[MemorySummary]]:
+        return separate_process_wrapper_fn(self._inference_memory, self.args.do_multi_processing)(*args, **kwargs)
+
+    def train_memory(self, *args, **kwargs) -> [Memory, Optional[MemorySummary]]:
+        return separate_process_wrapper_fn(self._train_memory, self.args.do_multi_processing)(*args, **kwargs)
+
     def run(self):
         result_dict = {model_name: {} for model_name in self.args.model_names}
         inference_result_time = copy.deepcopy(result_dict)
@@ -454,48 +656,66 @@ class Benchmark(ABC):
             train_result_time[model_name] = copy.deepcopy(model_dict)
             train_result_memory[model_name] = copy.deepcopy(model_dict)
 
+            inference_summary = train_summary = None
+
             for batch_size in self.args.batch_sizes:
                 for sequence_length in self.args.sequence_lengths:
                     if not self.args.no_inference:
                         if not self.args.no_memory:
-                            memory = self.inference(model_name, batch_size, sequence_length, trace_memory=True)
+                            memory, inference_summary = self.inference_memory(model_name, batch_size, sequence_length)
                             inference_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                         if not self.args.no_speed:
-                            time = self.inference(model_name, batch_size, sequence_length, trace_memory=False)
+                            time = self.inference_speed(model_name, batch_size, sequence_length)
                             inference_result_time[model_name]["result"][batch_size][sequence_length] = time
 
                     if self.args.training:
                         if not self.args.no_memory:
-                            memory = self.train(model_name, batch_size, sequence_length, trace_memory=True)
+                            memory, train_summary = self.train_memory(model_name, batch_size, sequence_length)
                             train_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                         if not self.args.no_speed:
-                            time = self.inference(model_name, batch_size, sequence_length, trace_memory=False)
+                            time = self.train_speed(model_name, batch_size, sequence_length)
                             train_result_time[model_name]["result"][batch_size][sequence_length] = time
 
         if not self.args.no_inference:
             if not self.args.no_speed:
-                self.print_fn("======= INFERENCE - SPEED - RESULT =======")
-                self.print_results(inference_result_time)
+                self.print_fn("\n" + 20 * "=" + ("INFERENCE - SPEED - RESULT").center(40) + 20 * "=")
+                self.print_results(inference_result_time, type_label="Time in s")
                 self.save_to_csv(inference_result_time, self.args.inference_time_csv_file)
+                if self.args.is_tpu:
+                    self.print_fn(
+                        "TPU was used for inference. Note that the time after compilation stabilized (after ~10 inferences model.forward(..) calls) was measured."
+                    )
 
             if not self.args.no_memory:
-                self.print_fn("======= INFERENCE - MEMORY - RESULT =======")
-                self.print_results(inference_result_memory)
+                self.print_fn("\n" + 20 * "=" + ("INFERENCE - MEMORY - RESULT").center(40) + 20 * "=")
+                self.print_results(inference_result_memory, type_label="Memory in MB")
                 self.save_to_csv(inference_result_memory, self.args.inference_memory_csv_file)
 
+            if self.args.trace_memory_line_by_line:
+                self.print_fn("\n" + 20 * "=" + ("INFERENCE - MEMOMRY - LINE BY LINE - SUMMARY").center(40) + 20 * "=")
+                self.print_memory_trace_statistics(inference_summary)
+
         if self.args.training:
             if not self.args.no_speed:
-                self.print_fn("======= TRAIN - SPEED - RESULT =======")
-                self.print_results(train_result_time)
+                self.print_fn("\n" + 20 * "=" + ("TRAIN - SPEED - RESULTS").center(40) + 20 * "=")
+                self.print_results(train_result_time, "Time in s")
                 self.save_to_csv(train_result_time, self.args.train_time_csv_file)
+                if self.args.is_tpu:
+                    self.print_fn(
+                        "TPU was used for training. Note that the time after compilation stabilized (after ~10 train loss=model.forward(...) + loss.backward() calls) was measured."
+                    )
 
             if not self.args.no_memory:
-                self.print_fn("======= TRAIN - MEMORY - RESULT =======")
-                self.print_results(train_result_memory)
+                self.print_fn("\n" + 20 * "=" + ("TRAIN - MEMORY - RESULTS").center(40) + 20 * "=")
+                self.print_results(train_result_memory, type_label="Memory in MB")
                 self.save_to_csv(train_result_memory, self.args.train_memory_csv_file)
 
+            if self.args.trace_memory_line_by_line:
+                self.print_fn("\n" + 20 * "=" + ("TRAIN - MEMOMRY - LINE BY LINE - SUMMARY").center(40) + 20 * "=")
+                self.print_memory_trace_statistics(train_summary)
+
         if not self.args.no_env_print:
-            self.print_fn("\n======== ENVIRONMENT - INFORMATION ========")
+            self.print_fn("\n" + 20 * "=" + ("ENVIRONMENT INFORMATION").center(40) + 20 * "=")
             self.print_fn(
                 "\n".join(["- {}: {}".format(prop, val) for prop, val in self.environment_info.items()]) + "\n"
             )
@@ -506,7 +726,14 @@ class Benchmark(ABC):
                 for key, value in self.environment_info.items():
                     writer.writerow([key, value])
 
-        return BenchmarkOutput(inference_result_time, inference_result_memory, train_result_time, train_result_memory)
+        return BenchmarkOutput(
+            inference_result_time,
+            inference_result_memory,
+            train_result_time,
+            train_result_memory,
+            inference_summary,
+            train_summary,
+        )
 
     @property
     def environment_info(self):
@@ -514,6 +741,11 @@ class Benchmark(ABC):
             info = {}
             info["transformers_version"] = version
             info["framework"] = self.framework
+            if self.framework == "PyTorch":
+                info["use_torchscript"] = self.args.torchscript
+            if self.framework == "TensorFlow":
+                info["eager_mode"] = self.args.eager_mode
+                info["use_xla"] = self.args.use_xla
             info["framework_version"] = self.framework_version
             info["python_version"] = platform.python_version()
             info["system"] = platform.system()
@@ -521,27 +753,31 @@ class Benchmark(ABC):
             info["architecture"] = platform.architecture()[0]
             info["date"] = datetime.date(datetime.now())
             info["time"] = datetime.time(datetime.now())
+            info["fp16"] = self.args.fp16
+            info["use_multiprocessing"] = self.args.do_multi_processing
+            info["only_pretrain_model"] = self.args.only_pretrain_model
 
-            try:
-                import psutil
-            except (ImportError):
+            if is_psutil_available():
+                info["cpu_ram_mb"] = bytes_to_mega_bytes(psutil.virtual_memory().total)
+            else:
                 logger.warning(
                     "Psutil not installed, we won't log available CPU memory."
                     "Install psutil (pip install psutil) to log available CPU memory."
                 )
                 info["cpu_ram_mb"] = "N/A"
-            else:
-                info["cpu_ram_mb"] = bytes_to_mega_bytes(psutil.virtual_memory().total)
 
-            info["use_gpu"] = self.is_gpu
-            if self.is_gpu:
-                info["num_gpus"] = self.args.n_gpu
-                try:
-                    from py3nvml import py3nvml
-
-                    py3nvml.nvmlInit()
-                    handle = py3nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
-                except ImportError:
+            info["use_gpu"] = self.args.is_gpu
+            if self.args.is_gpu:
+                info["num_gpus"] = 1  # TODO(PVP) Currently only single GPU is supported
+                if is_py3nvml_available():
+                    nvml.nvmlInit()
+                    handle = nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
+                    info["gpu"] = nvml.nvmlDeviceGetName(handle)
+                    info["gpu_ram_mb"] = bytes_to_mega_bytes(nvml.nvmlDeviceGetMemoryInfo(handle).total)
+                    info["gpu_power_watts"] = nvml.nvmlDeviceGetPowerManagementLimit(handle) / 1000
+                    info["gpu_performance_state"] = nvml.nvmlDeviceGetPerformanceState(handle)
+                    nvml.nvmlShutdown()
+                else:
                     logger.warning(
                         "py3nvml not installed, we won't log GPU memory usage. "
                         "Install py3nvml (pip install py3nvml) to log information about GPU."
@@ -550,37 +786,35 @@ class Benchmark(ABC):
                     info["gpu_ram_mb"] = "N/A"
                     info["gpu_power_watts"] = "N/A"
                     info["gpu_performance_state"] = "N/A"
-                except (OSError, py3nvml.NVMLError):
-                    logger.warning(
-                        "Error while initializing comunication with GPU. " "We won't log information about GPU."
-                    )
-                    info["gpu"] = "N/A"
-                    info["gpu_ram_mb"] = "N/A"
-                    info["gpu_power_watts"] = "N/A"
-                    info["gpu_performance_state"] = "N/A"
-                    py3nvml.nvmlShutdown()
-                else:
-                    info["gpu"] = py3nvml.nvmlDeviceGetName(handle)
-                    info["gpu_ram_mb"] = bytes_to_mega_bytes(py3nvml.nvmlDeviceGetMemoryInfo(handle).total)
-                    info["gpu_power_watts"] = py3nvml.nvmlDeviceGetPowerManagementLimit(handle) / 1000
-                    info["gpu_performance_state"] = py3nvml.nvmlDeviceGetPerformanceState(handle)
-                    py3nvml.nvmlShutdown()
+
+            info["use_tpu"] = self.args.is_tpu
+            # TODO(PVP): See if we can add more information about TPU
+            # see: https://github.com/pytorch/xla/issues/2180
 
             self._environment_info = info
         return self._environment_info
 
-    def print_results(self, result_dict):
+    def print_results(self, result_dict, type_label):
+        self.print_fn(80 * "-")
+        self.print_fn(
+            "Model Name".center(30) + "Batch Size".center(15) + "Seq Length".center(15) + type_label.center(15)
+        )
+        self.print_fn(80 * "-")
         for model_name in self.args.model_names:
-            self.print_fn("\t" + f"======= MODEL CHECKPOINT: {model_name} =======")
             for batch_size in result_dict[model_name]["bs"]:
                 for sequence_length in result_dict[model_name]["ss"]:
                     result = result_dict[model_name]["result"][batch_size][sequence_length]
                     if isinstance(result, float):
-                        self.print_fn(
-                            f"\t\t{model_name}/{batch_size}/{sequence_length}: " f"{(round(1000 * result) / 1000)}s"
-                        )
+                        result = round(1000 * result) / 1000
+                        result = "< 0.001" if result == 0.0 else str(result)
                     else:
-                        self.print_fn(f"\t\t{model_name}/{batch_size}/{sequence_length}: " f"{result} MB")
+                        result = str(result)
+                    self.print_fn(
+                        model_name[:30].center(30) + str(batch_size).center(15),
+                        str(sequence_length).center(15),
+                        result.center(15),
+                    )
+        self.print_fn(80 * "-")
 
     def print_memory_trace_statistics(self, summary: MemorySummary):
         self.print_fn(

src/transformers/file_utils.py

@@ -60,6 +60,35 @@ except (ImportError, AssertionError):
     _tf_available = False  # pylint: disable=invalid-name
 
 
+try:
+    import torch_xla.core.xla_model as xm  # noqa: F401
+
+    if _torch_available:
+        _torch_tpu_available = True  # pylint: disable=
+    else:
+        _torch_tpu_available = False
+except ImportError:
+    _torch_tpu_available = False
+
+
+try:
+    import psutil  # noqa: F401
+
+    _psutil_available = True
+
+except ImportError:
+    _psutil_available = False
+
+
+try:
+    import py3nvml  # noqa: F401
+
+    _py3nvml_available = True
+
+except ImportError:
+    _py3nvml_available = False
+
+
 try:
     from torch.hub import _get_torch_home
 
@@ -98,6 +127,18 @@ def is_tf_available():
     return _tf_available
 
 
+def is_torch_tpu_available():
+    return _torch_tpu_available
+
+
+def is_psutil_available():
+    return _psutil_available
+
+
+def is_py3nvml_available():
+    return _py3nvml_available
+
+
 def add_start_docstrings(*docstr):
     def docstring_decorator(fn):
         fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")