transformer上手（6）—— 微调预训练模型

1 加载数据集

以同义句判断任务为例（每次输入两个句子，判断它们是否为同义句），构建我们的第一个 Transformers 模型。我们选择蚂蚁金融语义相似度数据集 AFQMC 作为语料，它提供了官方的数据划分，训练集 / 验证集 / 测试集分别包含 34334 / 4316 / 3861 个句子对，标签 0 表示非同义句，1 表示同义句：

{"sentence1": "还款还清了，为什么花呗账单显示还要还款", "sentence2": "花呗全额还清怎么显示没有还款", "label": "1"}

1.1 Dataset

Pytorch 通过 Dataset 类和 DataLoader 类处理数据集和加载样本。同样地，这里我们首先继承 Dataset 类构造自定义数据集，以组织样本和标签。AFQMC 样本以 json 格式存储，因此我们使用 json 库按行读取样本，并且以行号作为索引构建数据集。

from torch.utils.data import Dataset
import jsonclass AFQMC(Dataset):def __init__(self, data_file):self.data = self.load_data(data_file)def load_data(self, data_file):Data = {}with open(data_file, 'rt') as f:for idx, line in enumerate(f):sample = json.loads(line.strip())Data[idx] = samplereturn Datadef __len__(self):return len(self.data)def __getitem__(self, idx):return self.data[idx]train_data = AFQMC('data/afqmc_public/train.json')
valid_data = AFQMC('data/afqmc_public/dev.json')print(train_data[0])# {'sentence1': '蚂蚁借呗等额还款可以换成先息后本吗', 'sentence2': '借呗有先息到期还本吗', 'label': '0'}

可以看到，我们编写的 AFQMC 类成功读取了数据集，每一个样本都以字典形式保存，分别以sentence1、sentence2 和 label 为键存储句子对和标签。

如果数据集非常巨大，难以一次性加载到内存中，我们也可以继承 IterableDataset 类构建迭代型数据集：

from torch.utils.data import IterableDataset
import jsonclass IterableAFQMC(IterableDataset):def __init__(self, data_file):self.data_file = data_filedef __iter__(self):with open(self.data_file, 'rt') as f:for line in f:sample = json.loads(line.strip())yield sampletrain_data = IterableAFQMC('data/afqmc_public/train.json')print(next(iter(train_data)))# {'sentence1': '蚂蚁借呗等额还款可以换成先息后本吗', 'sentence2': '借呗有先息到期还本吗', 'label': '0'}

1.2 DataLoader

接下来就需要通过 DataLoader 库按批 (batch) 加载数据，并且将样本转换成模型可以接受的输入格式。对于 NLP 任务，这个环节就是将每个 batch 中的文本按照预训练模型的格式进行编码（包括 Padding、截断等操作）。

通过手工编写 DataLoader 的批处理函数 collate_fn 来实现。首先加载分词器，然后对每个 batch 中的所有句子对进行编码，同时把标签转换为张量格式：

import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizercheckpoint = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)def collote_fn(batch_samples):batch_sentence_1, batch_sentence_2 = [], []batch_label = []for sample in batch_samples:batch_sentence_1.append(sample['sentence1'])batch_sentence_2.append(sample['sentence2'])batch_label.append(int(sample['label']))X = tokenizer(batch_sentence_1, batch_sentence_2, padding=True, truncation=True, return_tensors="pt")y = torch.tensor(batch_label)return X, ytrain_dataloader = DataLoader(train_data, batch_size=4, shuffle=True, collate_fn=collote_fn)batch_X, batch_y = next(iter(train_dataloader))
print('batch_X shape:', {k: v.shape for k, v in batch_X.items()})
print('batch_y shape:', batch_y.shape)
print(batch_X)
print(batch_y)batch_X shape: {'input_ids': torch.Size([4, 39]), 'token_type_ids': torch.Size([4, 39]), 'attention_mask': torch.Size([4, 39])
}
batch_y shape: torch.Size([4])# {'input_ids': tensor([
#         [ 101, 5709, 1446, 5543, 3118,  802,  736, 3952, 3952, 2767, 1408,  102,
#          3952, 2767, 1041,  966, 5543, 4500, 5709, 1446, 1408,  102,    0,    0,
#             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
#             0,    0,    0],
#         [ 101,  872, 8024, 2769, 6821, 5709, 1446, 4638, 7178, 6820, 2130,  749,
#          8024, 6929, 2582,  720, 1357, 3867,  749,  102, 1963, 3362, 1357, 3867,
#           749, 5709, 1446,  722, 1400, 8024, 1355, 4495, 4638, 6842, 3621, 2582,
#           720, 1215,  102],
#         [ 101, 1963,  862, 2990, 7770,  955, 1446,  677, 7361,  102, 6010, 6009,
#           955, 1446, 1963,  862,  712, 1220, 2990, 7583, 2428,  102,    0,    0,
#             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
#             0,    0,    0],
#         [ 101, 2582, 3416, 2990, 7770,  955, 1446, 7583, 2428,  102,  955, 1446,
#          2990, 4157, 7583, 2428, 1416,  102,    0,    0,    0,    0,    0,    0,
#             0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
#             0,    0,    0]]), 
#  'token_type_ids': tensor([
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,
#          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]), 
#  'attention_mask': tensor([
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
#          1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0,
#          0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])}
# tensor([1, 0, 1, 1])

可以看到，DataLoader 按照我们设置的 batch size 每次对 4 个样本进行编码，并且通过设置 padding=True 和 truncation=True 来自动对每个 batch 中的样本进行补全和截断。这里我们选择 BERT 模型作为 checkpoint，所以每个样本都被处理成了“$[CLS]sen1[SEP]sen2[SEP]$”的形式。

这种只在一个 batch 内进行补全的操作被称为动态补全 (Dynamic padding)，Hugging Face 也提供了 DataCollatorWithPadding 类来进行。

2 训练模型

2.1 构建模型

对于分类任务，可以直接使用我们前面介绍过的 AutoModelForSequenceClassification 类来完成。但是在实际操作中，除了使用预训练模型编码文本外，我们通常还会进行许多自定义操作，因此在大部分情况下我们都需要自己编写模型。

最简单的方式是首先利用 Transformers 库加载 BERT 模型，然后接一个全连接层完成分类：

from torch import nn
from transformers import AutoModeldevice = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device')class BertForPairwiseCLS(nn.Module):def __init__(self):super(BertForPairwiseCLS, self).__init__()self.bert_encoder = AutoModel.from_pretrained(checkpoint)self.dropout = nn.Dropout(0.1)self.classifier = nn.Linear(768, 2)def forward(self, x):bert_output = self.bert_encoder(**x)cls_vectors = bert_output.last_hidden_state[:, 0, :]cls_vectors = self.dropout(cls_vectors)logits = self.classifier(cls_vectors)return logitsmodel = BertForPairwiseCLS().to(device)
print(model)# Using cpu device
# NeuralNetwork(
#   (bert_encoder): BertModel(
#     (embeddings): BertEmbeddings(...)
#     (encoder): BertEncoder(...)
#     (pooler): BertPooler(
#       (dense): Linear(in_features=768, out_features=768, bias=True)
#       (activation): Tanh()
#     )
#   )
#   (dropout): Dropout(p=0.1, inplace=False)
#   (classifier): Linear(in_features=768, out_features=2, bias=True)
# )

这里模型首先将输入送入到 BERT 模型中，将每一个 token 都编码为维度为 768 的向量，然后从输出序列中取出第一个 [CLS] token 的编码表示作为整个句子对的语义表示，再送入到一个线性全连接层中预测两个类别的分数。这种方式简单粗暴，但是相当于在 Transformers 模型外又包了一层，因此无法再调用 Transformers 库预置的模型函数。

更为常见的写法是继承 Transformers 库中的预训练模型来创建自己的模型。例如这里我们可以继承 BERT 模型（BertPreTrainedModel 类）来创建一个与上面模型结构完全相同的分类器：

from torch import nn
from transformers import AutoConfig
from transformers import BertPreTrainedModel, BertModeldevice = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device')class BertForPairwiseCLS(BertPreTrainedModel):def __init__(self, config):super().__init__(config)self.bert = BertModel(config, add_pooling_layer=False)self.dropout = nn.Dropout(config.hidden_dropout_prob)self.classifier = nn.Linear(768, 2)self.post_init()def forward(self, x):bert_output = self.bert(**x)cls_vectors = bert_output.last_hidden_state[:, 0, :]cls_vectors = self.dropout(cls_vectors)logits = self.classifier(cls_vectors)return logitsconfig = AutoConfig.from_pretrained(checkpoint)
model = BertForPairwiseCLS.from_pretrained(checkpoint, config=config).to(device)
print(model)

注意，此时我们的模型是 Transformers 预训练模型的子类，因此需要通过预置的 from_pretrained 函数来加载模型参数。这种方式也使得我们可以更灵活地操作模型细节，例如这里 Dropout 层就可以直接加载 BERT 模型自带的参数值，而不用像上面一样手工赋值。

为了确保模型的输出符合我们的预期，我们尝试将一个 Batch 的数据送入模型：

outputs = model(batch_X)
print(outputs.shape)# torch.Size([4, 2])

可以看到模型输出了一个 $4\times 2$ 的张量，符合我们的预期（每个样本输出 2 维的 logits 值分别表示两个类别的预测分数，batch 内共 4 个样本）。

2.2 优化模型参数

正如之前介绍的那样，在训练模型时，我们将每一轮 Epoch 分为训练循环和验证/测试循环。在训练循环中计算损失、优化模型的参数，在验证/测试循环中评估模型的性能：

from tqdm.auto import tqdmdef train_loop(dataloader, model, loss_fn, optimizer, lr_scheduler, epoch, total_loss):progress_bar = tqdm(range(len(dataloader)))progress_bar.set_description(f'loss: {0:>7f}')finish_step_num = (epoch-1)*len(dataloader)model.train()for step, (X, y) in enumerate(dataloader, start=1):X, y = X.to(device), y.to(device)pred = model(X)loss = loss_fn(pred, y)optimizer.zero_grad()loss.backward()optimizer.step()lr_scheduler.step()total_loss += loss.item()progress_bar.set_description(f'loss: {total_loss/(finish_step_num + step):>7f}')progress_bar.update(1)return total_lossdef test_loop(dataloader, model, mode='Test'):assert mode in ['Valid', 'Test']size = len(dataloader.dataset)correct = 0model.eval()with torch.no_grad():for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)correct += (pred.argmax(1) == y).type(torch.float).sum().item()correct /= sizeprint(f"{mode} Accuracy: {(100*correct):>0.1f}%\n")

最后，将”训练循环”和”验证/测试循环”组合成 Epoch，就可以进行模型的训练和验证了。

与 Pytorch 类似，Transformers 库同样实现了很多的优化器，并且相比 Pytorch 固定学习率，Transformers 库的优化器会随着训练过程逐步减小学习率（通常会产生更好的效果）。例如我们前面使用过的 AdamW 优化器：

from transformers import AdamWoptimizer = AdamW(model.parameters(), lr=5e-5)

默认情况下，优化器会线性衰减学习率，对于上面的例子，学习率会线性地从 $5e-5$ 降到 0。为了正确地定义学习率调度器，我们需要知道总的训练步数 (step)，它等于训练轮数 (Epoch number) 乘以每一轮中的步数（也就是训练 dataloader 的大小）：

from transformers import get_schedulerepochs = 3
num_training_steps = epochs * len(train_dataloader)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=num_training_steps,
)
print(num_training_steps)# 25752

完整的训练过程如下：

from transformers import AdamW, get_schedulerlearning_rate = 1e-5
epoch_num = 3loss_fn = nn.CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=learning_rate)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=epoch_num*len(train_dataloader),
)total_loss = 0.
for t in range(epoch_num):print(f"Epoch {t+1}/{epoch_num}\n-------------------------------")total_loss = train_loop(train_dataloader, model, loss_fn, optimizer, lr_scheduler, t+1, total_loss)test_loop(valid_dataloader, model, mode='Valid')
print("Done!")# Using cuda device# Epoch 1/3
# -------------------------------
# loss: 0.552296: 100%|█████████| 8584/8584 [07:16<00:00, 19.65it/s]
# Valid Accuracy: 72.1%# Epoch 2/3
# -------------------------------
# loss: 0.501410: 100%|█████████| 8584/8584 [07:16<00:00, 19.66it/s]
# Valid Accuracy: 73.0%# Epoch 3/3
# -------------------------------
# loss: 0.450708: 100%|█████████| 8584/8584 [07:15<00:00, 19.70it/s]
# Valid Accuracy: 74.1%# Done!

2.3 保存和加载模型

在大多数情况下，我们还需要根据验证集上的表现来调整超参数以及选出最好的模型，最后再将选出的模型应用于测试集以评估性能。这里我们在测试循环时返回计算出的准确率，然后对上面的 Epoch 训练代码进行小幅的调整，以保存验证集上准确率最高的模型：

def test_loop(dataloader, model, mode='Test'):assert mode in ['Valid', 'Test']size = len(dataloader.dataset)correct = 0model.eval()with torch.no_grad():for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)correct += (pred.argmax(1) == y).type(torch.float).sum().item()correct /= sizeprint(f"{mode} Accuracy: {(100*correct):>0.1f}%\n")return correcttotal_loss = 0.
best_acc = 0.
for t in range(epoch_num):print(f"Epoch {t+1}/{epoch_num}\n-------------------------------")total_loss = train_loop(train_dataloader, model, loss_fn, optimizer, lr_scheduler, t+1, total_loss)valid_acc = test_loop(valid_dataloader, model, mode='Valid')if valid_acc > best_acc:best_acc = valid_accprint('saving new weights...\n')torch.save(model.state_dict(), f'epoch_{t+1}_valid_acc_{(100*valid_acc):0.1f}_model_weights.bin')
print("Done!")# Using cuda device# Epoch 1/3
# -------------------------------
# loss: 0.556518: 100%|█████████| 8584/8584 [07:51<00:00, 18.20it/s]
# Valid Accuracy: 71.8%# saving new weights...# Epoch 2/3
# -------------------------------
# loss: 0.506202: 100%|█████████| 8584/8584 [07:15<00:00, 19.71it/s]
# Valid Accuracy: 72.0%# saving new weights...# Epoch 3/3
# -------------------------------
# loss: 0.455851: 100%|█████████| 8584/8584 [07:16<00:00, 19.68it/s]
# Valid Accuracy: 74.1%# saving new weights...# Done!

可以看到，随着训练的进行，在验证集上的准确率逐步提升（71.8% -> 72.0% -> 74.1%）。因此，3 轮 Epoch 训练结束后，会在目录下保存下所有三轮模型的权重：

epoch_1_valid_acc_71.8_model_weights.bin
epoch_2_valid_acc_72.0_model_weights.bin
epoch_3_valid_acc_74.1_model_weights.bin

至此，我们手工构建的文本分类模型的训练过程就完成了，完整的训练代码如下：

import random
import os
import numpy as np
import json
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoConfig
from transformers import BertPreTrainedModel, BertModel
from transformers import AdamW, get_scheduler
from tqdm.auto import tqdmdef seed_everything(seed=1029):random.seed(seed)os.environ['PYTHONHASHSEED'] = str(seed)np.random.seed(seed)torch.manual_seed(seed)torch.cuda.manual_seed(seed)torch.cuda.manual_seed_all(seed)# some cudnn methods can be random even after fixing the seed# unless you tell it to be deterministictorch.backends.cudnn.deterministic = Truedevice = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device')
seed_everything(42)learning_rate = 1e-5
batch_size = 4
epoch_num = 3checkpoint = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)class AFQMC(Dataset):def __init__(self, data_file):self.data = self.load_data(data_file)def load_data(self, data_file):Data = {}with open(data_file, 'rt') as f:for idx, line in enumerate(f):sample = json.loads(line.strip())Data[idx] = samplereturn Datadef __len__(self):return len(self.data)def __getitem__(self, idx):return self.data[idx]train_data = AFQMC('data/afqmc_public/train.json')
valid_data = AFQMC('data/afqmc_public/dev.json')def collote_fn(batch_samples):batch_sentence_1, batch_sentence_2 = [], []batch_label = []for sample in batch_samples:batch_sentence_1.append(sample['sentence1'])batch_sentence_2.append(sample['sentence2'])batch_label.append(int(sample['label']))X = tokenizer(batch_sentence_1, batch_sentence_2, padding=True, truncation=True, return_tensors="pt")y = torch.tensor(batch_label)return X, ytrain_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, collate_fn=collote_fn)
valid_dataloader= DataLoader(valid_data, batch_size=batch_size, shuffle=False, collate_fn=collote_fn)class BertForPairwiseCLS(BertPreTrainedModel):def __init__(self, config):super().__init__(config)self.bert = BertModel(config, add_pooling_layer=False)self.dropout = nn.Dropout(config.hidden_dropout_prob)self.classifier = nn.Linear(768, 2)self.post_init()def forward(self, x):outputs = self.bert(**x)cls_vectors = outputs.last_hidden_state[:, 0, :]cls_vectors = self.dropout(cls_vectors)logits = self.classifier(cls_vectors)return logitsconfig = AutoConfig.from_pretrained(checkpoint)
model = BertForPairwiseCLS.from_pretrained(checkpoint, config=config).to(device)def train_loop(dataloader, model, loss_fn, optimizer, lr_scheduler, epoch, total_loss):progress_bar = tqdm(range(len(dataloader)))progress_bar.set_description(f'loss: {0:>7f}')finish_step_num = (epoch-1)*len(dataloader)model.train()for step, (X, y) in enumerate(dataloader, start=1):X, y = X.to(device), y.to(device)pred = model(X)loss = loss_fn(pred, y)optimizer.zero_grad()loss.backward()optimizer.step()lr_scheduler.step()total_loss += loss.item()progress_bar.set_description(f'loss: {total_loss/(finish_step_num + step):>7f}')progress_bar.update(1)return total_lossdef test_loop(dataloader, model, mode='Test'):assert mode in ['Valid', 'Test']size = len(dataloader.dataset)correct = 0model.eval()with torch.no_grad():for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)correct += (pred.argmax(1) == y).type(torch.float).sum().item()correct /= sizeprint(f"{mode} Accuracy: {(100*correct):>0.1f}%\n")return correctloss_fn = nn.CrossEntropyLoss()
optimizer = AdamW(model.parameters(), lr=learning_rate)
lr_scheduler = get_scheduler("linear",optimizer=optimizer,num_warmup_steps=0,num_training_steps=epoch_num*len(train_dataloader),
)total_loss = 0.
best_acc = 0.
for t in range(epoch_num):print(f"Epoch {t+1}/{epoch_num}\n-------------------------------")total_loss = train_loop(train_dataloader, model, loss_fn, optimizer, lr_scheduler, t+1, total_loss)valid_acc = test_loop(valid_dataloader, model, mode='Valid')if valid_acc > best_acc:best_acc = valid_accprint('saving new weights...\n')torch.save(model.state_dict(), f'epoch_{t+1}_valid_acc_{(100*valid_acc):0.1f}_model_weights.bin')
print("Done!")

这里我们还通过 seed_everything 函数手工设置训练过程中所有的随机数种子为 42，从而使得模型结果可以复现，而不是每次运行代码都是不同的结果。

最后，我们加载验证集上最优的模型权重，汇报其在测试集上的性能。由于 AFQMC 公布的测试集上并没有标签，无法评估性能，这里我们暂且用验证集代替进行演示：

model.load_state_dict(torch.load('epoch_3_valid_acc_74.1_model_weights.bin'))
test_loop(valid_dataloader, model, mode='Test')# Test Accuracy: 74.1%

注意：前面我们只保存了模型的权重，因此如果要单独调用上面的代码，需要首先实例化一个结构完全一样的模型，再通过 model.load_state_dict() 函数加载权重。

最终在测试集（这里用了验证集）上的准确率为 74.1%，与前面汇报的一致，也验证了模型参数的加载过程是正确的。