使用 Huggingface TFTrainer 类微调模型时如何指定损失函数？

Posted 2023-03-29

【中文标题】使用 Huggingface TFTrainer 类微调模型时如何指定损失函数？

【英文标题】：How to specify the loss function when finetuning a model using the Huggingface TFTrainer Class?

【发布时间】：2021-05-23 21:48:20

【问题描述】：

我遵循了下面给出的基本示例，来自：https://huggingface.co/transformers/training.html

from transformers import TFBertForSequenceClassification, TFTrainer, TFTrainingArguments

model = TFBertForSequenceClassification.from_pretrained("bert-large-uncased")

training_args = TFTrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total # of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=64,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
)

trainer = TFTrainer(
    model=model,                         # the instantiated ???? Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=tfds_train_dataset,    # tensorflow_datasets training dataset
    eval_dataset=tfds_test_dataset       # tensorflow_datasets evaluation dataset
)
trainer.train()

但似乎没有办法为分类器指定损失函数。 For-ex，如果我对二元分类问题进行微调，我会使用

tf.keras.losses.BinaryCrossentropy(from_logits=True)

否则我会使用

tf.keras.losses.CategoricalCrossentropy(from_logits=True)

我的设置如下：

transformers==4.3.2
tensorflow==2.3.1
python==3.6.12

【参考方案1】：

创建一个继承自 PreTrainedModel 的类，然后在它的前向函数中创建您各自的损失函数。

【讨论】：

你的建议工作正常 :) 但我喜欢一些最低限度运行的代码，以便它对未来的新手（也许还有我）有说明。

【参考方案2】：

Trainer 具有使用 compute_loss 的能力

有关更多信息，您可以查看文档：https://huggingface.co/docs/transformers/main_classes/trainer#:~:text=passed%20at%20init.-,compute_loss,-%2D%20Computes%20the%20loss

这是一个如何自定义 Trainer 以使用加权损失的示例（当您有不平衡的训练集时很有用）：

from torch import nn
from transformers import Trainer


class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # compute custom loss (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 2.0, 3.0]))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss