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

最新推荐文章于 2025-04-25 14:08:30 发布

页页读

最新推荐文章于 2025-04-25 14:08:30 发布

阅读量531

点赞数 4

分类专栏： transformers 大模型文章标签： transformers Trainer

本文链接：https://blog.csdn.net/u014386899/article/details/145597186

版权

问题描述

我在使用 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)

        # 计算 loss
        loss = self.loss_func(outputs, labels.float())

        if self.state