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

Question

我遵循了下面给出的基本示例，来自: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()

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

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

Answer 1

Trainer具有使用compute_loss的能力

https://huggingface.co/docs/transformers/main_classes/trainer#:~:text=passed%20at%20init.-，compute_loss，-%2D%20Computes%20the%20loss

下面是一个如何自定义训练器来使用加权损失的例子(当你有一个不平衡的训练集时很有用):

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

Answer 2

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