NLP（九）：pytorch用transformer库实现BERT

一、资源

（1）预训练模型权重

链接: https://pan.baidu.com/s/10BCm_qOlajUU3YyFDdLVBQ 密码: 1upi

（2）数据集选择的THUCNews，自行下载并整理出10w条数据，内容是10类新闻文本标题的中文分类问题（10分类），每类新闻标题数据量相等，为1w条。数据集可在我的百度网盘自行下载：链接: https://pan.baidu.com/s/1Crj4ELKtW8zRnNuaAkonPA 密码: p0wj。

（3）安装

pip install transformers

（4）参考：

https://zhuanlan.zhihu.com/p/112655246

https://spaces.ac.cn/archives/6736

二、简介

由于pytorch_pretrained_bert库是transformers的老版库，不再进行更新了。所以以下对原文章代码进行了更新，换成以transformers为框架的代码，并且将打印输出设置的更加简约。本文章适用于初学者，有兴趣的可上手尝试。

----------------- 分割线 ----------------

之前用bert一直都是根据keras-bert封装库操作的，操作非常简便（可参考苏剑林大佬博客当Bert遇上Keras：这可能是Bert最简单的打开姿势），这次想要来尝试一下基于pytorch的bert实践。

最近pytorch大火，而目前很少有博客完整的给出基于pytorch的bert的应用代码，本文从最简单的中文文本分类入手，一步一步的给出每段代码~ （代码简单清晰，读者有兴趣可上手实践）

（1）首先安装transformers库， 即：pip install transformers（版本为4.4.2）；

（2）然后下载预训练模型权重，这里下载的是 chinese_roberta_wwm_ext_pytorch ，下载链接为中文BERT-wwm系列模型 （这里可选择多种模型），如果下载不了，可在我的百度网盘下载：链接: https://pan.baidu.com/s/10BCm_qOlajUU3YyFDdLVBQ 密码: 1upi；

（3）数据集选择的THUCNews，自行下载并整理出10w条数据，内容是10类新闻文本标题的中文分类问题（10分类），每类新闻标题数据量相等，为1w条。数据集可在我的百度网盘自行下载：链接: https://pan.baidu.com/s/1Crj4ELKtW8zRnNuaAkonPA 密码: p0wj。

下图为数据集展示（最后10行），格式为"title \t label"，标题和所属类别两列用'\t'分隔。

废话少说，下面进入代码阶段。（训练环境为Google Colab，GPU为T4，显存大约15G）

1 导入必要的库

import pandas as pd 
import numpy as np 
import json, time 
from tqdm import tqdm 
from sklearn.metrics import accuracy_score, classification_report
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
from transformers import BertModel, BertConfig, BertTokenizer, AdamW, get_cosine_schedule_with_warmup
import warnings
warnings.filterwarnings('ignore')


bert_path = "bert_model/"    # 该文件夹下存放三个文件（'vocab.txt', 'pytorch_model.bin', 'config.json'）
tokenizer = BertTokenizer.from_pretrained(bert_path)   # 初始化分词器

2 预处理数据集

input_ids, input_masks, input_types,  = [], [], []  # input char ids, segment type ids,  attention mask
labels = []         # 标签
maxlen = 30      # 取30即可覆盖99%
 
with open("news_title_dataset.csv", encoding='utf-8') as f:
    for i, line in tqdm(enumerate(f)): 
        title, y = line.strip().split('\t')

        # encode_plus会输出一个字典，分别为'input_ids', 'token_type_ids', 'attention_mask'对应的编码
        # 根据参数会短则补齐，长则切断
        encode_dict = tokenizer.encode_plus(text=title, max_length=maxlen, 
                                            padding='max_length', truncation=True)
        
        input_ids.append(encode_dict['input_ids'])
        input_types.append(encode_dict['token_type_ids'])
        input_masks.append(encode_dict['attention_mask'])

        labels.append(int(y))

input_ids, input_types, input_masks = np.array(input_ids), np.array(input_types), np.array(input_masks)
labels = np.array(labels)
print(input_ids.shape, input_types.shape, input_masks.shape, labels.shape)

输出：（27秒，速度较快）
100000it [00:27, 3592.75it/s]
(100000, 30) (100000, 30) (100000, 30) (100000,)

3 切分训练集、验证集和测试集

# 随机打乱索引
idxes = np.arange(input_ids.shape[0])
np.random.seed(2019)   # 固定种子
np.random.shuffle(idxes)
print(idxes.shape, idxes[:10])


# 8:1:1 划分训练集、验证集、测试集
input_ids_train, input_ids_valid, input_ids_test = input_ids[idxes[:80000]], input_ids[idxes[80000:90000]], input_ids[idxes[90000:]]
input_masks_train, input_masks_valid, input_masks_test = input_masks[idxes[:80000]], input_masks[idxes[80000:90000]], input_masks[idxes[90000:]] 
input_types_train, input_types_valid, input_types_test = input_types[idxes[:80000]], input_types[idxes[80000:90000]], input_types[idxes[90000:]]

y_train, y_valid, y_test = labels[idxes[:80000]], labels[idxes[80000:90000]], labels[idxes[90000:]]

print(input_ids_train.shape, y_train.shape, input_ids_valid.shape, y_valid.shape, 
      input_ids_test.shape, y_test.shape)

输出：

4加载到高效的DataLoader

BATCH_SIZE = 64  # 如果会出现OOM问题，减小它
# 训练集
train_data = TensorDataset(torch.LongTensor(input_ids_train), 
                           torch.LongTensor(input_masks_train), 
                           torch.LongTensor(input_types_train), 
                           torch.LongTensor(y_train))
train_sampler = RandomSampler(train_data)  
train_loader = DataLoader(train_data, sampler=train_sampler, batch_size=BATCH_SIZE)

# 验证集
valid_data = TensorDataset(torch.LongTensor(input_ids_valid), 
                          torch.LongTensor(input_masks_valid),
                          torch.LongTensor(input_types_valid), 
                          torch.LongTensor(y_valid))
valid_sampler = SequentialSampler(valid_data)
valid_loader = DataLoader(valid_data, sampler=valid_sampler, batch_size=BATCH_SIZE)

# 测试集（是没有标签的）
test_data = TensorDataset(torch.LongTensor(input_ids_test), 
                          torch.LongTensor(input_masks_test),
                          torch.LongTensor(input_types_test))
test_sampler = SequentialSampler(test_data)
test_loader = DataLoader(test_data, sampler=test_sampler, batch_size=BATCH_SIZE)

5 定义bert模型

# 定义model
class Bert_Model(nn.Module):
    def __init__(self, bert_path, classes=10):
        super(Bert_Model, self).__init__()
        self.config = BertConfig.from_pretrained(bert_path)
        self.bert = BertModel.from_pretrained(bert_path)
        self.fc = nn.Linear(self.config.hidden_size, classes)  # 直接分类
        
        
    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        outputs = self.bert(input_ids, attention_mask, token_type_ids)
        out_pool = outputs[1]   # 池化后的输出 
        logit = self.fc(out_pool)
        return logit

可以发现，bert模型的定义由于高效简易的封装库存在，使得定义模型较为容易，如果想要在bert之后加入cnn/rnn等层，可在这里定义。

6 实例化bert模型

def get_parameter_number(model):
    #  打印模型参数
    total_num = sum(p.numel() for p in model.parameters())
    trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return 'Total parameters: {}, Trainable parameters: {}'.format(total_num, trainable_num)

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCHS = 5
model = Bert_Model(bert_path).to(DEVICE)
print(get_parameter_number(model))

输出：Total parameters: 102275338, Trainable parameters: 102275338

7 定义优化器

optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=1e-4) #AdamW优化器
scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=len(train_loader),
                                            num_training_steps=EPOCHS*len(train_loader))
# 学习率先线性warmup一个epoch，然后cosine式下降。

8 定义训练函数和验证测试函数

# 评估模型性能，在验证集上
def evaluate(model, data_loader, device):
    model.eval()
    val_true, val_pred = [], []
    with torch.no_grad():
        for idx, (ids, att, tpe, y) in (enumerate(data_loader)):
            y_pred = model(ids.to(device), att.to(device), tpe.to(device))
            y_pred = torch.argmax(y_pred, dim=1).detach().cpu().numpy().tolist()
            val_pred.extend(y_pred)
            val_true.extend(y.squeeze().cpu().numpy().tolist())
    
    return accuracy_score(val_true, val_pred)  #返回accuracy


# 测试集没有标签，需要预测提交
def predict(model, data_loader, device):
    model.eval()
    val_pred = []
    with torch.no_grad():
        for idx, (ids, att, tpe) in tqdm(enumerate(data_loader)):
            y_pred = model(ids.to(device), att.to(device), tpe.to(device))
            y_pred = torch.argmax(y_pred, dim=1).detach().cpu().numpy().tolist()
            val_pred.extend(y_pred)
    return val_pred


def train_and_eval(model, train_loader, valid_loader, 
                   optimizer, scheduler, device, epoch):
    best_acc = 0.0
    patience = 0
    criterion = nn.CrossEntropyLoss()
    for i in range(epoch):
        """训练模型"""
        start = time.time()
        model.train()
        print("***** Running training epoch {} *****".format(i+1))
        train_loss_sum = 0.0
        for idx, (ids, att, tpe, y) in enumerate(train_loader):
            ids, att, tpe, y = ids.to(device), att.to(device), tpe.to(device), y.to(device)  
            y_pred = model(ids, att, tpe)
            loss = criterion(y_pred, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()   # 学习率变化
            
            train_loss_sum += loss.item()
            if (idx + 1) % (len(train_loader)//5) == 0:    # 只打印五次结果
                print("Epoch {:04d} | Step {:04d}/{:04d} | Loss {:.4f} | Time {:.4f}".format(
                          i+1, idx+1, len(train_loader), train_loss_sum/(idx+1), time.time() - start))
                # print("Learning rate = {}".format(optimizer.state_dict()['param_groups'][0]['lr']))

        """验证模型"""
        model.eval()
        acc = evaluate(model, valid_loader, device)  # 验证模型的性能
        if acc > best_acc:
            best_acc = acc
            torch.save(model.state_dict(), "best_bert_model.pth") 
        
        print("current acc is {:.4f}, best acc is {:.4f}".format(acc, best_acc))
        print("time costed = {}s \n".format(round(time.time() - start, 5)))

9 开始训练和验证模型

# 训练和验证评估
train_and_eval(model, train_loader, valid_loader, optimizer, scheduler, DEVICE, EPOCHS)

输出：（训练时间较长，500s左右一个epoch，这里只训练了2个epoch，验证集便得到了0.9680的accuracy）

10 加载最优模型进行测试

# 加载最优权重对测试集测试
model.load_state_dict(torch.load("best_bert_model.pth"))
pred_test = predict(model, test_loader, DEVICE)
print("\n Test Accuracy = {} \n".format(accuracy_score(y_test, pred_test)))
print(classification_report(y_test, pred_test, digits=4))

输出：测试集准确率为96.72%

------------------------------------

经过以上10步，即可建立起较为完整的基于pytorch的bert文本分类体系，代码也较为简单易懂，对读者有帮助记得点个赞支持一下呀~

-完结-