SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

QnA example: add speed metric (#20522)

Browse Source
This commit is contained in:
Wang, Yi
2022-12-02 01:04:19 +08:00
committed by GitHub
parent b67ac44296
commit d752337baa
1 changed files with 25 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
examples/pytorch/question-answering/trainer_qa.py
View File

@@ -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 @@ class QuestionAnsweringTrainer(Trainer):
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 @@ class QuestionAnsweringTrainer(Trainer):
)
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 @@ class QuestionAnsweringTrainer(Trainer):
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 @@ class QuestionAnsweringTrainer(Trainer):
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 @@ class QuestionAnsweringTrainer(Trainer):
)
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 @@ class QuestionAnsweringTrainer(Trainer):
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)