蒸馏后的模型进行evaluate，报错AxisError: axis 2 is out of bounds for array of dimension 1

[MrRace]

使用 https://github.com/airaria/TextBrewer/blob/master/examples/notebook_examples/msra_ner.ipynb
在蒸馏之后，对蒸馏后的模型进行evaluate:

trainer.evaluate()

报错：

AxisError                                 Traceback (most recent call last)
File <timed eval>:1

File /opt/conda/lib/python3.8/site-packages/transformers/trainer.py:2811, in Trainer.evaluate(self, eval_dataset, ignore_keys, metric_key_prefix)
   2808 start_time = time.time()
   2810 eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
-> 2811 output = eval_loop(
   2812     eval_dataloader,
   2813     description="Evaluation",
   2814     # No point gathering the predictions if there are no metrics, otherwise we defer to
   2815     # self.args.prediction_loss_only
   2816     prediction_loss_only=True if self.compute_metrics is None else None,
   2817     ignore_keys=ignore_keys,
   2818     metric_key_prefix=metric_key_prefix,
   2819 )
   2821 total_batch_size = self.args.eval_batch_size * self.args.world_size
   2822 output.metrics.update(
   2823     speed_metrics(
   2824         metric_key_prefix,
   (...)
   2828     )
   2829 )

File /opt/conda/lib/python3.8/site-packages/transformers/trainer.py:3096, in Trainer.evaluation_loop(self, dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix)
   3092         metrics = self.compute_metrics(
   3093             EvalPrediction(predictions=all_preds, label_ids=all_labels, inputs=all_inputs)
   3094         )
   3095     else:
-> 3096         metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
   3097 else:
   3098     metrics = {}

Cell In[61], line 11, in compute_metrics(p)
      9     print("predictions 0 shape=", predictions[0].shape) # (10, 374, 7)
     10     print("predictions 1 shape=", len(predictions[1]))
---> 11     predictions = np.argmax(predictions, axis=2)
     12 #     print("2-predictions=", predictions)
     13     
     14     # Remove ignored index (special tokens)
     15     print("Here!!!")

File <__array_function__ internals>:180, in argmax(*args, **kwargs)

File /opt/conda/lib/python3.8/site-packages/numpy/core/fromnumeric.py:1216, in argmax(a, axis, out, keepdims)
   1129 """
   1130 Returns the indices of the maximum values along an axis.
   1131 
   (...)
   1213 (2, 1, 4)
   1214 """
   1215 kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {}
-> 1216 return _wrapfunc(a, 'argmax', axis=axis, out=out, **kwds)

File /opt/conda/lib/python3.8/site-packages/numpy/core/fromnumeric.py:54, in _wrapfunc(obj, method, *args, **kwds)
     52 bound = getattr(obj, method, None)
     53 if bound is None:
---> 54     return _wrapit(obj, method, *args, **kwds)
     56 try:
     57     return bound(*args, **kwds)

File /opt/conda/lib/python3.8/site-packages/numpy/core/fromnumeric.py:43, in _wrapit(obj, method, *args, **kwds)
     41 except AttributeError:
     42     wrap = None
---> 43 result = getattr(asarray(obj), method)(*args, **kwds)
     44 if wrap:
     45     if not isinstance(result, mu.ndarray):

AxisError: axis 2 is out of bounds for array of dimension 1

根据报错信息，似乎是用于train阶段的compute_metrics不能用于evaluate:

def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)

    # Remove ignored index (special tokens)
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    
    results = metric.compute(predictions=true_predictions, references=true_labels)
    return {
        "precision": results["overall_precision"],
        "recall": results["overall_recall"],
        "f1": results["overall_f1"],
        "accuracy": results["overall_accuracy"],
    }

[MrRace]

@airaria Could you help to solve it？Thanks a lot!

[MrRace]

should change raw predictions = np.argmax(predictions, axis=2) to predictions = np.argmax(predictions[0], axis=2) when do eval

[airaria]

thank a lot!
The notebook is a bit outdated ...I'll fix this.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[stale]

Closing the issue, since no updates observed. Feel free to re-open if you need any further assistance.