微调BERT时调用auroc()报错：缺少必填参数'task'求解决方案

我在包含1.5k条数据的小型数据集上微调BERT base uncased，执行训练代码：

trainer.fit(model, data_module)

自定义模型代码如下：

class ElectionTagger(pl.LightningModule):
  def __init__(self, n_classes: int, n_training_steps=None, n_warmup_steps=None):
    super().__init__()
    self.bert = BertModel.from_pretrained(BERT_MODEL_NAME, return_dict=True)
    self.classifier = nn.Linear(self.bert.config.hidden_size, n_classes)
    self.n_training_steps = n_training_steps
    self.n_warmup_steps = n_warmup_steps
    self.criterion = nn.BCELoss()

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.bert(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)
    loss = 0
    if labels is not None:
        loss = self.criterion(output, labels)
    return loss, output

  def training_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True, logger=True)
    return {"loss": loss, "predictions": outputs, "labels": labels}

  def validation_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("val_loss", loss, prog_bar=True, logger=True)
    return loss

  def test_step(self, batch, batch_idx):
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]
    labels = batch["labels"]
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("test_loss", loss, prog_bar=True, logger=True)
    return loss

  def training_epoch_end(self, outputs):
    labels = []
    predictions = []
    for output in outputs:
      for out_labels in output["labels"].detach().cpu():
        labels.append(out_labels)
      for out_predictions in output["predictions"].detach().cpu():
        predictions.append(out_predictions)
    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)
    for i, name in enumerate(LABEL_COLUMNS):
      class_roc_auc = auroc(predictions[:, i], labels[:, i]) ##### ERROR ARISES HERE###
      self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", class_roc_auc, self.current_epoch)

  def configure_optimizers(self):
    optimizer = AdamW(self.parameters(), lr=2e-5)
    scheduler = get_linear_schedule_with_warmup(
      optimizer,
      num_warmup_steps=self.n_warmup_steps,
      num_training_steps=self.n_training_steps
    )
    return dict(
      optimizer=optimizer,
      lr_scheduler=dict(
        scheduler=scheduler,
        interval='step'
      )
    )

运行时触发错误：

TypeError: auroc() missing 1 required positional argument: 'task'

解决方案 #

该错误源于新版torchmetrics库中的auroc函数强制要求指定task参数，用于明确任务类型。结合你使用BCELoss和sigmoid输出的多标签二分类场景，提供以下三种解决方式：

1. 直接添加task参数 #

在调用auroc时显式指定task="binary"，因为每个标签都是独立的二分类任务：

class_roc_auc = auroc(predictions[:, i], labels[:, i], task="binary")

2. 使用多标签专属AUROC指标 #

针对多标签场景，推荐使用MultilabelAUROC指标类，更贴合业务逻辑：

• 首先导入指标：

from torchmetrics import MultilabelAUROC

• 修改training_epoch_end方法：

def training_epoch_end(self, outputs):
    labels = []
    predictions = []
    for output in outputs:
      labels.extend(output["labels"].detach().cpu())
      predictions.extend(output["predictions"].detach().cpu())
    labels = torch.stack(labels).int()
    predictions = torch.stack(predictions)
    
    # 计算整体多标签AUROC
    multilabel_auroc = MultilabelAUROC(num_labels=len(LABEL_COLUMNS), task="multilabel")
    overall_auroc = multilabel_auroc(predictions, labels)
    self.logger.experiment.add_scalar("overall_roc_auc/Train", overall_auroc, self.current_epoch)
    
    # 单独计算每个标签的AUROC
    for i, name in enumerate(LABEL_COLUMNS):
        class_auroc = auroc(predictions[:, i], labels[:, i], task="binary")
        self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", class_auroc, self.current_epoch)

3. 回退到旧版torchmetrics（不推荐） #

若不想修改代码，可安装v0.11.x版本的torchmetrics，该版本的auroc无需task参数：

pip install torchmetrics==0.11.4

注意：此方法不推荐，旧版本可能存在未修复的bug，且无法享受后续功能更新。

内容的提问来源于stack exchange，提问作者Minor Project