Dataset适配#

在 Transformers实战——使用Trainer类训练和评估自己的数据和模型中，使用的数据集类别是from datasets import Dataset，但在常用的实现中，pytorch项目会使用torch.utils.data 中的数据集。

为了让pytorch中的Dataset类适配transformers中的Trainer ，需要对类别中的def __getitem__(self, idx) 方法进行修改，如下：

1
    def __getitem__(self, idx):
2
        # 具体实现，可以自由发挥
3
        text = self.texts[idx]
4
        label = self.labels[idx]
5
        encoding = self.tokenizer.encode_plus(
6
            text,
7
            add_special_tokens=True,
8
            return_token_type_ids=False,
9
            truncation=True,
10
            return_attention_mask=True,
11
            return_tensors='pt',
12
        )
13

14
        # 下面的返回要改成字典形式
15
        return {
16
            'input_ids': encoding['input_ids'].flatten(),
17
            'attention_mask': encoding['attention_mask'].flatten(),
18
            'labels': torch.tensor(label, dtype=torch.long)
19
        }

简单来说就是把返回改成字典形式，且字典的键（key）需要与模型的forward对应。

DataCollator#

上面的改动可以控制每一个样本的返回值，但模型训练时的样本往往是一批一批的。如果每个样本的输入形状（shape）固定不变的话，到上一步模型就能正常运行了。

但有些时候每个样本的输入形状不一样，如文本分类（每篇文章字数不同），就会导致训练不能正常进行。

这种情况一般有两种方法（以文本分类为例）：

修改__getitem__ ：添加一个最大长度 max_length 把每个文本扩展（0填充）或截取到固定长度（如512）。
使用transformes中的DataCollator 类，如文本填充常用的DataCollatorWithPadding 。这种方式不用修改Dataset类的源码。

使用DataCollator 的样例如下：

1
from transformers import DataCollatorWithPadding
2

3
# 这里的tokenizer变量可以传入BertTokenizer类
4
data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors='pt')
5

6
    trainer = Trainer(
7
        model=model,  # the instantiated 🤗 Transformers model to be trained 需要训练的模型
8
        args=training_args,  # training arguments, defined above 训练参数
9
        train_dataset=train_dataset,  # training dataset 训练集
10
        eval_dataset=dev_dataset,  # evaluation dataset 测试集
11
        compute_metrics=compute_metrics,  # 计算指标方法
12
        ########
13
        data_collator=data_collator, # 传入DataCollator
14
        ########
15
    )

这种方法可以动态地根据每一批次的最大文本长度进行补全，可以一定程度节省内存消耗。

原文链接：CSDN