【transformers.Trainer填坑】在自定义compute_metrics时logits和labels数据维度不一致问题

问题描述

我在使用 transformers.Trainer 训练我的模型时，我自定义了 compute_loss 函数和compute_metrics函数，我的模型是一个简单的二分类模型。

在自定义 compute_loss 时这样写的：

def compute_loss(self, model, inputs, return_outputs=False):"""重写 Trainer.compute_loss:1) 提取字典中的 images, bboxes, locs, labels 等2) 用 vision_encoder 先处理图像，得到特征3) 用下游 model 做预测4) 计算并返回 loss"""# 前向传播outputs, labels = model(**inputs)  # (bz, num_classes), or (bz*num_frames, num_classes)batch_size = inputs['labels'].shape[0]outputs = outputs.squeeze()  # (bz*num_frames)if batch_size == 1:outputs = outputs.unsqueeze(0)# 计算 lossloss = self.loss_func(outputs, labels.float())if self.state.global_step % 10 == 0 and self.state.global_step > 0:# 以50个step为间隔打印pred_probs = torch.sigmoid(outputs)preds = (pred_probs > 0.5).int()logger.info(f"[global_step={self.state.global_step}] preds={preds.tolist()} / labels={labels.tolist()} / loss={loss.item():.4f}")# compute metricaccuracy = accuracy_score(labels.cpu().numpy(), preds.cpu().numpy())precision = precision_score(labels.cpu().numpy(), preds.cpu().numpy())recall = recall_score(labels.cpu().numpy(), preds.cpu().numpy())logger.info(f"[global_step={self.state.global_step}] accuracy={accuracy:.4f} / precision={precision:.4f} / recall={recall:.4f}")# 返回 (loss, outputs) 或者只返回 lossreturn (loss, outputs) if return_outputs else loss

于是就出现了报错，像这样的：

File "/opt/conda/lib/python3.9/site-packages/transformers/trainer.py", line 3754, in predictoutput = eval_loop(File "/opt/conda/lib/python3.9/site-packages/transformers/trainer.py", line 3966, in evaluation_loopmetrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))File "/workspace/train/object_query/train.py", line 281, in compute_metricscorrect_num = preds == labels
ValueError: operands could not be broadcast together with shapes (11720,) (12104,)output = eval_loop(File "/opt/conda/lib/python3.9/site-packages/transformers/trainer.py", line 3966, in evaluation_loopmetrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))File "/workspace/train/object_query/train.py", line 281, in compute_metricscorrect_num = preds == labels
ValueError: operands could not be broadcast together with shapes (11720,) (12104,)

原因

该问题是 transformers.Trainer 内部有一段对outputs的操作造成的：

if isinstance(outputs, dict):logits = tuple(v for k, v in outputs.items() if k not in ignore_keys + ["loss"])
else:logits = outputs[1:]

这里当 outputs 不是字典时，会把第一个位置的元素offset掉。

解决

Refer to here
所以，我们应该在返回那里这样写：

return (loss, {"label": outputs}) if return_outputs else loss