有时候我们并不想使用Transformers来训练别人的预训练模型，而是想用来训练自己的模型，并且不想写训练过程代码。这时，我们可以按照一定的要求定义数据集和模型，就可以使用Trainer类来直接训练和评估模型，不需要写那些训练步骤了。

使用Trainer类训练自己模型步骤如下，以网络退化现象和残差网络效果中的残差模型为例：

导入必要库#

1
import torch
2
from torch import nn
3
from datasets import Dataset
4
from transformers import Trainer, TrainingArguments
5

6
# 驱动选择
7
device = "cuda" if torch.cuda.is_available() else "cpu"

准备数据#

1
X = torch.zeros((26, 26), dtype=torch.float32).to(device=device)
2
labels = []
3
for i in range(26):
4
    labels.append((i+1) % 26)
5
    X[i][i] = 1.
6
labels = torch.tensor(labels)
7
dataset = Dataset.from_dict({'x':X, 'labels':labels})

注意

Trainer训练时，会将dataset中的数据按照对应的键值传入，因此需要在自己模型的forward方法中接收键值变量。如上例，需要将方法写为：forward(self, x, labels)

构建网络#

1
# 残差网络
2
class RN(nn.Module):
3
    def __init__(self):
4
        super(RN, self).__init__()
5
        self.linear_stack = nn.Sequential(
6
            nn.Linear(26, 64),
7
            nn.Hardsigmoid(),
8
            nn.Linear(64, 26),
9
            nn.Hardsigmoid(),
10
        )
11

12
        self.linear_stack_2 = nn.Sequential(
13
            nn.Linear(26, 64),
14
            nn.Hardsigmoid(),
15
            nn.Linear(64, 64),
16
            nn.Hardsigmoid(),
17
        )
18

19
        self.output_layer = nn.Linear(64, 26)
20

21
        self.loss_f = nn.CrossEntropyLoss()
22

23
    def forward(self, x, labels, mode='train'):
24
        y = self.linear_stack(x)
25
        # 残差
26
        y = y+x
27
        y = self.linear_stack_2(y)
28
        y = self.output_layer(y)
29

30
        if mode is 'train':
31
            return {
32
                'loss':self.loss_f(y, labels),
33
                'predictions':y
34
            }
35

36
        return y

网络与之前完全一致，只是改了下前向传播方法，这里的mode=’train’是为了后续自己使用模型加的，也可以不要。

创建Trainer类#

1
# 生成模型实例
2
model = RN().to(device=device)
3

4
def compute_metrics(pred):
5
    labels = pred.label_ids
6
    preds = pred.predictions.argmax(-1)
7
    acc = (labels == preds).sum()/len(labels)
8
    return {
9
        'accuracy': acc,
10
    }
11

12
training_args = TrainingArguments(
13
    output_dir='./results',         # output directory 结果输出地址
14
    num_train_epochs=1000,          # total # of training epochs 训练总批次
15
    per_device_train_batch_size=1,  # batch size per device during training 训练批大小
16
    per_device_eval_batch_size=1,   # batch size for evaluation 评估批大小
17
    logging_dir='./logs/rn_log',    # directory for storing logs 日志存储位置
18
    learning_rate=1e-3,             # 学习率
19
    save_steps=False,               # 不保存检查点
20
)
21

22
trainer = Trainer(
23
    model=model,                      # the instantiated 🤗 Transformers model to be trained 需要训练的模型
24
    args=training_args,               # training arguments, defined above 训练参数
25
    train_dataset=dataset,            # training dataset 训练集
26
    eval_dataset=dataset,             # evaluation dataset 测试集
27
    compute_metrics=compute_metrics   # 计算指标方法
28
)
29

30
trainer.train()
31
trainer.evaluate()

训练过程：

请添加图片描述

评估结果：

请添加图片描述

保存模型#

1
trainer.save_model()

加载并使用模型#

1
checkpoint = torch.load('./results/pytorch_model.bin')
2
model = RN().to(device)
3
model.load_state_dict(checkpoint)
4
model(X.to(device), torch.tensor(labels).to(device))['predictions'].argmax(1)

请添加图片描述

原文链接：CSDN