Fix trainer evaluation (#4363)

* fix loss calculation in evaluation

* fix evaluation on TPU when prediction_loss_only is True
This commit is contained in:
Suraj Patil 2020-05-15 00:09:44 +05:30 committed by GitHub
parent 035678efdb
commit 2d05480174
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -679,7 +679,7 @@ class Trainer:
model.eval()
for inputs in tqdm(dataloader, desc=description):
has_labels = any(inputs.get(k) is not None for k in ["labels", "masked_lm_labels"])
has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])
for k, v in inputs.items():
inputs[k] = v.to(self.args.device)
@ -703,7 +703,7 @@ class Trainer:
else:
label_ids = np.append(label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
if is_tpu_available():
if is_tpu_available() and preds is not None and label_ids is not None:
# tpu-comment: Get all predictions and labels from all worker shards of eval dataset
preds = xm.mesh_reduce("eval_preds", preds, np.concatenate)
label_ids = xm.mesh_reduce("eval_out_label_ids", label_ids, np.concatenate)