Repository: huggingface/transformers
Pull request: QnA example: add speed metric

Examples:
@sgugger


Changed files:
filename: examples/pytorch/question-answering/trainer_qa.py
patch: @@ -15,9 +15,11 @@
 """
 A subclass of `Trainer` specific to Question-Answering tasks
 """
+import math
+import time
 
 from transformers import Trainer, is_torch_tpu_available
-from transformers.trainer_utils import PredictionOutput
+from transformers.trainer_utils import PredictionOutput, speed_metrics
 
 
 if is_torch_tpu_available(check_device=False):
@@ -40,6 +42,7 @@ def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metr
         compute_metrics = self.compute_metrics
         self.compute_metrics = None
         eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        start_time = time.time()
         try:
             output = eval_loop(
                 eval_dataloader,
@@ -51,7 +54,15 @@ def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metr
             )
         finally:
             self.compute_metrics = compute_metrics
-
+        total_batch_size = self.args.eval_batch_size * self.args.world_size
+        output.metrics.update(
+            speed_metrics(
+                metric_key_prefix,
+                start_time,
+                num_samples=output.num_samples,
+                num_steps=math.ceil(output.num_samples / total_batch_size),
+            )
+        )
         if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
             # Only the main node write the results by default
             eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)
@@ -61,6 +72,7 @@ def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metr
             for key in list(metrics.keys()):
                 if not key.startswith(f"{metric_key_prefix}_"):
                     metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
+            metrics.update(output.metrics)
         else:
             metrics = {}
 
@@ -82,6 +94,7 @@ def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_ke
         compute_metrics = self.compute_metrics
         self.compute_metrics = None
         eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
+        start_time = time.time()
         try:
             output = eval_loop(
                 predict_dataloader,
@@ -93,6 +106,15 @@ def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_ke
             )
         finally:
             self.compute_metrics = compute_metrics
+        total_batch_size = self.args.eval_batch_size * self.args.world_size
+        output.metrics.update(
+            speed_metrics(
+                metric_key_prefix,
+                start_time,
+                num_samples=output.num_samples,
+                num_steps=math.ceil(output.num_samples / total_batch_size),
+            )
+        )
 
         if self.post_process_function is None or self.compute_metrics is None:
             return output
@@ -104,5 +126,5 @@ def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_ke
         for key in list(metrics.keys()):
             if not key.startswith(f"{metric_key_prefix}_"):
                 metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
-
+        metrics.update(output.metrics)
         return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)