基于PyTorch的BERT中文分类模型

2020-09-25

本文利用了transformers中的BertModel，对部分cnews数据集进行了文本分类，用来对BERT模型练手还是不错的。

数据描述

数据集是从清华大学的THUCNews中提取出来的部分数据。

训练集中有5万条数据，分成了10类，每类5000条数据。

1	{'体育': 5000, '娱乐': 5000, '家居': 5000, '房产': 5000, '教育': 5000, '时尚': 5000, '时政': 5000, '游戏': 5000, '科技': 5000, '财经': 5000}

验证集中有5000条数据，每类500条数据。

1	{'体育': 500, '娱乐': 500, '家居': 500, '房产': 500, '教育': 500, '时尚': 500, '时政': 500, '游戏': 500, '科技': 500, '财经': 500}

模型描述

整个分类模型首先把句子输入到Bert预训练模型，然后将句子的embedding输入给一个全连接层，最后把全连接层的输出输入到softmax中。模型代码如下：

# model.py
import torch
import torch.nn as nn
from transformers import  BertModel

class BertClassifier(nn.Module):
    def __init__(self, bert_config):
        super().__init__()
        self.bert = BertModel(config=bert_config)
        self.classifier = nn.Linear(bert_config.hidden_size, bert_config.num_labels)
        self.softmax = nn.functional.softmax

    def forward(self, input_ids, attention_mask, token_type_ids):
        bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        pooled = bert_output[1]
        logits = self.classifier(pooled)
        return self.softmax(logits,1)

数据处理

生成BERT的输入，并封装成TensorDataset

def get_bert_input(text, tokenizer, max_len=512):
    '''
        生成单个句子的BERT模型的三个输入
        参数:   
            text: 文本(单个句子)
            tokenizer: 分词器
            max_len: 文本分词后的最大长度
        返回值:
            input_id, attention_mask, token_type_id
    '''
    cls_token = '[CLS]'
    sep_token = '[SEP]'

    word_piece_list = tokenizer.tokenize(text)  #分词
    input_id = tokenizer.convert_tokens_to_ids(word_piece_list) #把分词结果转成id
    if len(input_id) > max_len-2:   #如果input_id的长度大于max_len，则进行截断操作
        input_id = input_id[:510]
    input_id = tokenizer.build_inputs_with_special_tokens(input_id) #对input_id补上[CLS]、[SEP]

    attention_mask = [] # 注意力的mask，把padding部分给遮蔽掉
    for i in range(len(input_id)):
        attention_mask.append(1)    # 句子的原始部分补1
    while len(attention_mask) < max_len:
        attention_mask.append(0)    # padding部分补0

    while len(input_id) < max_len:  # 如果句子长度小于max_len, 做padding，在句子后面补0
        input_id.append(0)

    token_type_id = [0] * max_len # 第一个句子为0，第二个句子为1，第三个句子为0 ..., 也可记为segment_id

    assert len(input_id) == len(token_type_id) == len(attention_mask)

    return input_id, attention_mask, token_type_id
    

def load_data(filename):
    '''
    读取数据文件，将数据封装起来返回，便于BERT输入
    输入：
        数据文件名
    输出：
        已经处理好的适合BERT输入的dataset
    '''
    with open(filename, 'r') as rf:
        data = rf.readlines()
    
    label_ids = []
    input_ids = []
    attention_mask = []
    token_type_ids = []

    tokenizer = BertTokenizer(vocab_file='../chinese-bert_chinese_wwm_pytorch/vocab.txt')

    for item in data:
        label, text = item.strip().split('\t')
        label_id = labels.index(label)
        input_id, attention_mask_id, token_type_id = get_bert_input(text, tokenizer)

        input_ids.append(input_id)
        attention_mask.append(attention_mask_id)
        token_type_ids.append(token_type_id)
        label_ids.append(label_id)
    
    input_ids = torch.tensor([i for i in input_ids], dtype=torch.long)
    attention_mask = torch.tensor([i for i in attention_mask], dtype=torch.long)
    token_type_ids = torch.tensor([i for i in token_type_ids], dtype=torch.long)
    label_ids = torch.tensor([i for i in label_ids], dtype=torch.long)
    dataset = TensorDataset(input_ids, attention_mask, token_type_ids, label_ids)
    return dataset

训练

训练模型，并对验证集进行验证，然后把最优模型保存下来

def main():

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 获取到dataset
    train_data = load_data('cnews/cnews.train.txt')
    valid_data = load_data('cnews/cnews.val.txt')
    #test_data = load_data('cnews/cnews.test.txt')

    batch_size = 4
    # 生成Batch
    train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
    valid_dataloader = DataLoader(valid_data, batch_size=batch_size, shuffle=False)
    #test_dataloader = DataLoader(valid_data, batch_size=batch_size, shuffle=False)

    # 读取BERT的配置文件
    bert_config = BertConfig.from_pretrained('../chinese-bert_chinese_wwm_pytorch')
    bert_config.num_labels = num_labels

    # 初始化模型
    model = BertClassifier(bert_config)
    model.to(device)

    # 参数设置
    EPOCHS = 10
    learning_rate = 5e-6    #Learning Rate不宜太大
    optimizer = AdamW(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()

    with open('output.txt', 'w') as wf:
        wf.write('Batch Size: ' + str(batch_size) + '\tLearning Rate: ' + str(learning_rate) + '\n')

    best_acc = 0
    # 开始训练 
    for Epoch in range(1, EPOCHS+1):
        losses = 0      # 损失
        accuracy = 0    # 准确率
        print('Epoch:',Epoch)

        model.train()
        for batch_index, batch in enumerate(train_dataloader):
            input_ids = batch[0].to(device)
            attention_mask = batch[1].to(device)
            token_type_ids = batch[2].to(device)
            label_ids = batch[3].to(device)

            output = model(    # forward
                input_ids=input_ids, 
                attention_mask=attention_mask, 
                token_type_ids=token_type_ids, 
            )

            loss = criterion(output, label_ids)
            losses += loss.item()

            pred_labels = torch.argmax(output, dim=1)   # 预测出的label
            acc = torch.sum(pred_labels == label_ids.to(device)).item() / len(pred_labels) #acc
            accuracy += acc

            print('Epoch: %d ｜ Train: | Batch: %d / %d | Acc: %f | Loss: %f' % (Epoch, batch_index+1, len(train_dataloader), acc, loss.item()))

            model.zero_grad()
            loss.backward()
            optimizer.step()


        average_loss = losses / len(train_dataloader)
        average_acc = accuracy / len(train_dataloader)

        print('\tTrain ACC:', average_acc, '\tLoss:', average_loss)
        with open('output.txt', 'a') as rf:
            output_to_file = '\nEpoch: ' + str(Epoch) + '\tTrain ACC:' + str(average_acc) + '\tLoss: ' + str(average_loss)
            rf.write(output_to_file)
        
        # 验证
        model.eval()
        losses = 0      # 损失
        accuracy = 0    # 准确率
        for batch_index, batch in enumerate(valid_dataloader):
            input_ids = batch[0].to(device)
            attention_mask = batch[1].to(device)
            token_type_ids = batch[2].to(device)
            label_ids = batch[3].to(device)

            output = model(    # forward
                input_ids=input_ids, 
                attention_mask=attention_mask, 
                token_type_ids=token_type_ids, 
            )
            
            loss = criterion(output, label_ids)
            losses += loss.item()

            pred_labels = torch.argmax(output, dim=1)   # 预测出的label
            acc = torch.sum(pred_labels == label_ids.to(device)).item() / len(pred_labels) #acc
            accuracy += acc

        average_loss = losses / len(valid_dataloader)
        average_acc = accuracy / len(valid_dataloader)

        print('\tValid ACC:', average_acc, '\tLoss:', average_loss)
        with open('output.txt', 'a') as rf:
            output_to_file = '\nEpoch: ' + str(Epoch) + '\tValid ACC:' + str(average_acc) + '\tLoss: ' + str(average_loss) + '\n'
            rf.write(output_to_file)
        
        if average_acc > best_acc:
            best_acc = average_acc
            torch.save(model.state_dict(), 'models/best_model.pkl')

预测

用训练好的最优模型对新闻进行预测

import torch
from model import BertClassifier
from transformers import BertTokenizer, BertConfig
from train import get_bert_input

labels = ['体育', '娱乐', '家居', '房产', '教育', '时尚', '时政', '游戏', '科技', '财经']
bert_config = BertConfig.from_pretrained('../chinese-bert_chinese_wwm_pytorch')
bert_config.num_labels = len(labels)
model = BertClassifier(bert_config)
model.load_state_dict(torch.load('models/best_model.pkl', map_location=torch.device('cpu')))

tokenizer = BertTokenizer(vocab_file='../chinese-bert_chinese_wwm_pytorch/vocab.txt')

print('新闻类别分类')
while True:
    text = input('Input: ')
    input_id, attention_mask, token_type_id = get_bert_input(text, tokenizer)

    input_id = torch.tensor([input_id], dtype=torch.long)
    attention_mask = torch.tensor([attention_mask], dtype=torch.long)
    token_type_id = torch.tensor([token_type_id], dtype=torch.long)

    predicted = model(
        input_id,
        attention_mask,
        token_type_id,
    )
    pred_label = torch.argmax(predicted, dim=1)

    print('Label:', labels[pred_label])

结果

# output.txt

Batch Size: 4   Learning Rate: 5e-06

Epoch: 1    Train ACC:0.68856   Loss: 1.7757255261802674
Epoch: 1    Valid ACC:0.8158    Loss: 1.6466600142478942

Epoch: 2    Train ACC:0.89862   Loss: 1.56370301820755
Epoch: 2    Valid ACC:0.837 Loss: 1.623342474269867

Epoch: 3    Train ACC:0.91918   Loss: 1.5420743488025666
Epoch: 3    Valid ACC:0.8672    Loss: 1.5936675718307496

Epoch: 4    Train ACC:0.92298   Loss: 1.5383553547668456
Epoch: 4    Valid ACC:0.859 Loss: 1.6021957096099853

Epoch: 5    Train ACC:0.92944   Loss: 1.5316217193889619
Epoch: 5    Valid ACC:0.8744    Loss: 1.5861366262435914

Epoch: 6    Train ACC:0.93062   Loss: 1.5303730747890472
Epoch: 6    Valid ACC:0.8856    Loss: 1.5743153604507447

Epoch: 7    Train ACC:0.93456   Loss: 1.5266000279617309
Epoch: 7    Valid ACC:0.8758    Loss: 1.5852036945343018

Epoch: 8    Train ACC:0.92862   Loss: 1.5324853926753998
Epoch: 8    Valid ACC:0.8166    Loss: 1.64386196641922

Epoch: 9    Train ACC:0.93212   Loss: 1.5289581029891968
Epoch: 9    Valid ACC:0.8546    Loss: 1.6062824898719787

Epoch: 10   Train ACC:0.9339    Loss: 1.5273041619777679
Epoch: 10   Valid ACC:0.8514    Loss: 1.6099982411384584

完整项目放在了GitHub上: https://github.com/illiterate/BertClassifier