NLP（三十五）使用keras-bert实现文本多分类任务_如何用bert在5000条数据集中训练模型分类

  本文将会介绍如何使用keras-bert实现文本多分类任务，其中对BERT进行微调。

项目结构

其中依赖的Python第三方模块如下：

pandas==0.23.4
Keras==2.3.1
keras_bert==0.83.0
numpy==1.16.4
1
2
3
4

数据集

  本文采用的多分类数据集为sougou小分类数据集和THUCNews数据集，简介如下：

• sougou小分类数据集

共有5个类别，分别为体育、健康、军事、教育、汽车。划分为训练集和测试集，其中训练集每个分类800条样本，测试集每个分类100条样本。

• THUCNews数据集

共有10个分类，类别为：体育, 财经, 房产, 家居, 教育, 科技, 时尚, 时政, 游戏, 娱乐。数据集划分为：训练集: 5000 * 10，测试集: 1000 * 10。

模型训练

  模型训练脚本model_train.py的完整代码如下：

# -*- coding: utf-8 -*-
import json
import codecs
import pandas as pd
import numpy as np
from keras_bert import load_trained_model_from_checkpoint, Tokenizer
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam

# 建议长度<=510
maxlen = 300
BATCH_SIZE = 8
config_path = './chinese_L-12_H-768_A-12/bert_config.json'
checkpoint_path = './chinese_L-12_H-768_A-12/bert_model.ckpt'
dict_path = './chinese_L-12_H-768_A-12/vocab.txt'


token_dict = {}
with codecs.open(dict_path, 'r', 'utf-8') as reader:
    for line in reader:
        token = line.strip()
        token_dict[token] = len(token_dict)


class OurTokenizer(Tokenizer):
    def _tokenize(self, text):
        R = []
        for c in text:
            if c in self._token_dict:
                R.append(c)
            else:
                R.append('[UNK]')   # 剩余的字符是[UNK]
        return R


tokenizer = OurTokenizer(token_dict)


def seq_padding(X, padding=0):
    L = [len(x) for x in X]
    ML = max(L)
    return np.array([
        np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X
    ])


class DataGenerator:

    def __init__(self, data, batch_size=BATCH_SIZE):
        self.data = data
        self.batch_size = batch_size
        self.steps = len(self.data) // self.batch_size
        if len(self.data) % self.batch_size != 0:
            self.steps += 1

    def __len__(self):
        return self.steps

    def __iter__(self):
        while True:
            idxs = list(range(len(self.data)))
            np.random.shuffle(idxs)
            X1, X2, Y = [], [], []
            for i in idxs:
                d = self.data[i]
                text = d[0][:maxlen]
                x1, x2 = tokenizer.encode(first=text)
                y = d[1]
                X1.append(x1)
                X2.append(x2)
                Y.append(y)
                if len(X1) == self.batch_size or i == idxs[-1]:
                    X1 = seq_padding(X1)
                    X2 = seq_padding(X2)
                    Y = seq_padding(Y)
                    yield [X1, X2], Y
                    [X1, X2, Y] = [], [], []


# 构建模型
def create_cls_model(num_labels):
    bert_model = load_trained_model_from_checkpoint(config_path, checkpoint_path, seq_len=None)

    for layer in bert_model.layers:
        layer.trainable = True

    x1_in = Input(shape=(None,))
    x2_in = Input(shape=(None,))

    x = bert_model([x1_in, x2_in])
    cls_layer = Lambda(lambda x: x[:, 0])(x)    # 取出[CLS]对应的向量用来做分类
    p = Dense(num_labels, activation='softmax')(cls_layer)     # 多分类

    model = Model([x1_in, x2_in], p)
    model.compile(
        loss='categorical_crossentropy',
        optimizer=Adam(1e-5),   # 用足够小的学习率
        metrics=['accuracy']
    )
    # model.summary()

    return model


if __name__ == '__main__':

    # 数据处理, 读取训练集和测试集
    print("begin data processing...")
    train_df = pd.read_csv("data/cnews/cnews_train.csv").fillna(value="")
    test_df = pd.read_csv("data/cnews/cnews_test.csv").fillna(value="")

    labels = train_df["label"].unique()
    with open("label.json", "w", encoding="utf-8") as f:
        f.write(json.dumps(dict(zip(range(len(labels)), labels)), ensure_ascii=False, indent=2))

    train_data = []
    test_data = []
    for i in range(train_df.shape[0]):
        label, content = train_df.iloc[i, :]
        label_id = [0] * len(labels)
        for j, _ in enumerate(labels):
            if _ == label:
                label_id[j] = 1
        train_data.append((content, label_id))

    for i in range(test_df.shape[0]):
        label, content = test_df.iloc[i, :]
        label_id = [0] * len(labels)
        for j, _ in enumerate(labels):
            if _ == label:
                label_id[j] = 1
        test_data.append((content, label_id))

    print("finish data processing!")

    # 模型训练
    model = create_cls_model(len(labels))
    train_D = DataGenerator(train_data)
    test_D = DataGenerator(test_data)

    print("begin model training...")
    model.fit_generator(
        train_D.__iter__(),
        steps_per_epoch=len(train_D),
        epochs=3,
        validation_data=test_D.__iter__(),
        validation_steps=len(test_D)
    )

    print("finish model training!")

    # 模型保存
    model.save('cls_cnews.h5')
    print("Model saved!")

    result = model.evaluate_generator(test_D.__iter__(), steps=len(test_D))
    print("模型评估结果:", result)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

其中模型结构如下：

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            (None, None)         0                                            
__________________________________________________________________________________________________
input_2 (InputLayer)            (None, None)         0                                            
__________________________________________________________________________________________________
model_2 (Model)                 (None, None, 768)    101677056   input_1[0][0]                    
                                                                 input_2[0][0]                    
__________________________________________________________________________________________________
lambda_1 (Lambda)               (None, 768)          0           model_2[1][0]                    
__________________________________________________________________________________________________
dense_1 (Dense)                 (None, 10)           7690        lambda_1[0][0]                   
==================================================================================================
Total params: 101,684,746
Trainable params: 101,684,746
Non-trainable params: 0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

在上述模型中，我们取取出[CLS]对应的向量，后接全连接层，激活函数采用Softmax函数，就完成多分类模型的搭建了，非常简单方便。

模型评估

  模型评估脚本model_evaluate.py的完整代码如下：

# -*- coding: utf-8 -*-
# 模型评估脚本
import json
import numpy as np
import pandas as pd
from keras.models import load_model
from keras_bert import get_custom_objects
from sklearn.metrics import classification_report

from model_train import token_dict, OurTokenizer

maxlen = 300

# 加载训练好的模型
model = load_model("cls_cnews.h5", custom_objects=get_custom_objects())
tokenizer = OurTokenizer(token_dict)
with open("label.json", "r", encoding="utf-8") as f:
    label_dict = json.loads(f.read())


# 对单句话进行预测
def predict_single_text(text):
    # 利用BERT进行tokenize
    text = text[:maxlen]
    x1, x2 = tokenizer.encode(first=text)
    X1 = x1 + [0] * (maxlen - len(x1)) if len(x1) < maxlen else x1
    X2 = x2 + [0] * (maxlen - len(x2)) if len(x2) < maxlen else x2

    # 模型预测并输出预测结果
    predicted = model.predict([[X1], [X2]])
    y = np.argmax(predicted[0])
    return label_dict[str(y)]


# 模型评估
def evaluate():
    test_df = pd.read_csv("data/cnews/cnews_test.csv").fillna(value="")
    true_y_list, pred_y_list = [], []
    for i in range(test_df.shape[0]):
        print("predict %d samples" % (i+1))
        true_y, content = test_df.iloc[i, :]
        pred_y = predict_single_text(content)
        true_y_list.append(true_y)
        pred_y_list.append(pred_y)

    return classification_report(true_y_list, pred_y_list, digits=4)


output_data = evaluate()
print("model evaluate result:\n")
print(output_data)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

运行上述代码，对两个数据集进行评估，结果如下：

• sougou数据集

模型参数: batch_size = 8, maxlen = 256, epoch=10

评估结果:

                  precision    recall  f1-score   support

          体育     0.9802    1.0000    0.9900        99
          健康     0.9495    0.9495    0.9495        99
          军事     1.0000    1.0000    1.0000        99
          教育     0.9307    0.9495    0.9400        99
          汽车     0.9895    0.9495    0.9691        99

    accuracy                         0.9697       495
   macro avg     0.9700    0.9697    0.9697       495
weighted avg     0.9700    0.9697    0.9697       495
1
2
3
4
5
6
7
8
9
10
11

• THUCNews数据集

模型参数: batch_size = 8, maxlen = 300, epoch=3

评估结果:

                precision    recall  f1-score   support

          体育     0.9970    0.9990    0.9980      1000
          娱乐     0.9890    0.9890    0.9890      1000
          家居     0.9949    0.7820    0.8757      1000
          房产     0.8006    0.8710    0.8343      1000
          教育     0.9753    0.9480    0.9615      1000
          时尚     0.9708    0.9980    0.9842      1000
          时政     0.9318    0.9560    0.9437      1000
          游戏     0.9851    0.9950    0.9900      1000
          科技     0.9689    0.9970    0.9828      1000
          财经     0.9377    0.9930    0.9645      1000

    accuracy                         0.9528     10000
   macro avg     0.9551    0.9528    0.9524     10000
weighted avg     0.9551    0.9528    0.9524     10000
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

模型预测

  模型预测脚本model_predict.py的完整代码如下：

# -*- coding: utf-8 -*-
# @Time : 2020/12/23 15:28
# @Author : Jclian91
# @File : model_predict.py
# @Place : Yangpu, Shanghai
# 模型预测脚本

import time
import json
import numpy as np

from model_train import token_dict, OurTokenizer
from keras.models import load_model
from keras_bert import get_custom_objects

maxlen = 256

# 加载训练好的模型
model = load_model("cls_sougou.h5", custom_objects=get_custom_objects())
tokenizer = OurTokenizer(token_dict)
with open("label.json", "r", encoding="utf-8") as f:
    label_dict = json.loads(f.read())

s_time = time.time()
# 预测示例语句
text = "说到硬派越野SUV，你会想起哪些车型？是被称为“霸道”的丰田 普拉多 (配置 | 询价) ，还是被叫做“山猫”的帕杰罗，亦或者是“渣男专车”奔驰大G、" \
       "“沙漠王子”途乐。总之，随着世界各国越来越重视对环境的保护，那些大排量的越野SUV在不久的将来也会渐渐消失在我们的视线之中，所以与其错过，" \
       "不如趁着还年轻，在有生之年里赶紧去入手一台能让你心仪的硬派越野SUV。而今天我想要来跟你们聊的，正是全球公认的十大硬派越野SUV，" \
       "越野迷们看完之后也不妨思考一下，到底哪款才是你的菜，下面话不多说，赶紧开始吧。"


# 利用BERT进行tokenize
text = text[:maxlen]
x1, x2 = tokenizer.encode(first=text)

X1 = x1 + [0] * (maxlen-len(x1)) if len(x1) < maxlen else x1
X2 = x2 + [0] * (maxlen-len(x2)) if len(x2) < maxlen else x2

# 模型预测并输出预测结果
predicted = model.predict([[X1], [X2]])
y = np.argmax(predicted[0])


print("原文: %s" % text)
print("预测标签: %s" % label_dict[str(y)])
e_time = time.time()
print("cost time:", e_time-s_time)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

  我们在新的样本上进行模型预测。

• sougou数据集

原文: 说到硬派越野SUV，你会想起哪些车型？是被称为“霸道”的丰田 普拉多 (配置 | 询价) ，还是被叫做“山猫”的帕杰罗，亦或者是“渣男专车”奔驰大G、“沙漠王子”途乐。总之，随着世界各国越来越重视对环境的保护，那些大排量的越野SUV在不久的将来也会渐渐消失在我们的视线之中，所以与其错过，不如趁着还年轻，在有生之年里赶紧去入手一台能让你心仪的硬派越野SUV。而今天我想要来跟你们聊的，正是全球公认的十大硬派越野SUV，越野迷们看完之后也不妨思考一下，到底哪款才是你的菜，下面话不多说，赶紧开始吧。
预测标签: 汽车

原文: 【#美30架战机在阿拉斯加海岸大象漫步#】据美国艾尔森空军基地网站消息称，近日美国空军30架战斗机和2架空中加油机自艾尔森空军基地起飞，在阿拉斯加海岸完成了”大象漫步“式的演习。
预测标签: 军事

原文: “十三五”期间，我国义务教育三科统编教材实现所有年级全覆盖；普通高中三科统编教材已覆盖20个省份，预计2022年前实现所有省份全覆盖，2025年实现所有年级全覆盖。昨日，在教育部新闻发布会上，教育部教材局局长田慧生透露，义务教育课程方案和各学科课程标准修订明年完成。
预测标签: 教育
1
2
3
4
5
6
7
8

• THUCNews数据集

原文: 北京时间12月26日,2020-21赛季NBA圣诞大战如约上演。在一场焦点对决中，洛杉矶湖人在主场与达拉斯独行侠遭遇。全场打完，湖人138-115轻取独行侠，拿到赛季首胜，同时也送给对手2连败。
预测标签: 体育

原文: 近两年来，手机屏幕就开始不断升级，高刷新率也成为一种趋势，就算性能做的再好，手机屏幕不能流畅真实的展现出来，也会很大程度上影响使用感受，所以屏幕就是手机硬件的窗户，可以预见未来的高端手机在冲击性能的同时，必然会对提高对屏幕的要求。
预测标签: 科技

原文: 松江佘山板块已太久没有豪宅入市，令区域内不少高端置业客寂寞难耐。不过好在，不久前火爆登场的国贸佘山原墅一下子满足了这类客户的需求。  无论是社区规划、产品户型、装修配置，还是设计手法，相较于周边千万级别墅，国贸佘山原墅也是不逞多让，更令人惊喜的是别墅的品质却仅需公寓的价格！
预测标签: 房产
1
2
3
4
5
6
7
8

总结

  本项目已经开源，Github地址为： https://github.com/percent4/keras_bert_text_classification 。
  后续将会介绍如何使用keras-bert实现文本多标签分类任务。
  2020年12月26日于上海浦东