使用带有变形金刚拥抱面的权重

Question

我遇到了这个使用 Longformer 执行文本分类的教程。 我遇到了这两个链接—— 一个和两个讨论在数据不平衡时使用 class 权重。

# instantiate the trainer class and check for available devices
trainer = Trainer(
    model=model,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=train_data,
    eval_dataset=test_data
)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
device

我不确定如何修改上面的代码以包含 class 权重，如下所示（从上面的最后一个链接复制的代码）

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 1

您现在可以简单地拥有一个 CustomTrainer，而不是实例化Trainer CustomTrainer 。 这意味着，而不是写

trainer = Trainer(...)

你只需写

trainer = CustomTrainer(...)

由于您从Trainer class 继承了所有功能和默认行为，因此唯一的区别是自定义compute_loss() function。