使用keras-bert进行中文文本分类+Google colab运行源码

前文介绍了BERT的原理。在实际应用中，BERT要比其理论本身要简单的多。这里我们利用Github的中文BERT预训练的结果（地址），进行实际的文档分类。

数据集

为了便于进行比较，文档分类的数据集来自Github的这个地址。该数据集使用了THUCNews的一个子集，使用了其中的10个分类，每个分类6500条数据。
类别如下：
体育, 财经, 房产, 家居, 教育, 科技, 时尚, 时政, 游戏, 娱乐
数据集划分如下：

• 训练集: 5000*10
• 验证集: 500*10
  Github上该数据集的地址在百度网盘上，可以用multicloud直接将百度网盘的文件传到google drive上。

基本结构

利用BERT的预训练模型进行文档分类，其原理如下：

其原理是在BERT的Transform Layer的最高一层的第一个输出，添加Dense+Softmax对文档进行分类，loss函数取cross entropy函数。因此，真正的核心代码也就如下几行：

inputs = bert_model.inputs[:2]
x=bert_model.layers[-1].output
x=keras.layers.Lambda(lambda x: x[:, 0], name='slice')(x)
x=keras.layers.Dense(units=3072, activation=keras.backend.tanh)(x)
x=keras.layers.Dropout(rate=0.1,seed=2019)(x)
x=keras.layers.Dense(units=num_labels, activation=keras.backend.softmax)(x)
model=keras.Model(inputs, x)
model.compile(optimizer=keras.optimizers.Adam(1e-4),
            loss=keras.losses.sparse_categorical_crossentropy,
            metrics=[keras.metrics.sparse_categorical_accuracy])

预训练的模型采用了小参量的模型RBT3，RBT3实际上只有3个transform layer，而不是完整的12个transform layer。因此，只能算是“低配穷人版”的BERT，但从后面的结果来看表现已经相当不错了。

框架

为了实现上述功能，本文采用了Keras框架。Keras框架的好处是糙快猛，便于深度学习的入门者快速上手，而针对BERT模型，也已经有开源的keras-bert可以直接拿来使用。

平台

源代码在Google colab平台上，使用Google免费提供的GPU跑通。本来本人想试TPU的，但是发现依据keras-bert建议的TPU训练代码，模型无法收敛，因此还有待牛人进一步解决了…

代码

接下来直接上代码。

下载数据到Colab目录

首先将RBT3对应的google drive的地址，直接add到自己的google drive账号里，这样就可以不用下载，直接在colab上将文件解压就可以。

!cp drive/My\ Drive/chinese_rbt3_L-3_H-768_A-12.zip .
!mkdir cnews
!mkdir model
!cp drive/My\ Drive/cnews/* cnews/
!rm -rf model
!unzip chinese_rbt3_L-3_H-768_A-12.zip -d model/

输出

Archive:  chinese_rbt3_L-3_H-768_A-12.zip
  inflating: model/bert_config_rbt3.json  
  inflating: model/bert_model.ckpt.data-00000-of-00001  
  inflating: model/__MACOSX/._bert_model.ckpt.data-00000-of-00001  
  inflating: model/bert_model.ckpt.index  
  inflating: model/__MACOSX/._bert_model.ckpt.index  
  inflating: model/bert_model.ckpt.meta  
  inflating: model/__MACOSX/._bert_model.ckpt.meta  
  inflating: model/vocab.txt         
  inflating: model/__MACOSX/._vocab.txt  

安装库

!pip install --upgrade --force-reinstall keras-bert keras-rectified-adam

导入模块

import numpy as np
import os
import keras
# from tensorflow.python import keras
from keras_bert import load_trained_model_from_checkpoint #用于加载预训练的bert
from keras_bert import get_pretrained, PretrainedList, get_checkpoint_paths
from keras_bert import Tokenizer #用于对输入的文章“分词”
import codecs
from keras.preprocessing.sequence import pad_sequences

注意到keras-bert引用了两种keras库：

• tensorflow.python.keras:
• keras:
  keras-bert默认使用的是keras库，因此模型也要对应import keras。如果在python里执行

os.environ['TF_KERAS'] = '1'

那么接下来模型要import tensorflow.python.keras。如果弄混了，在build model的时候就会有千奇百怪的错误。
如果采用colab的TPU，就需要用tensorflow.python.keras包了，因为在keras本身对TPU支持目前还有限，对TPU的操作主要依赖于tensorflow，而且还必须是tensorflow 1.x的版本，tensorflow 2.x版本目前对TPU的支持不好。目前colab默认的还是tensorflow 1.x的版本。

配置数据的路径

# 与训练的bert model的路径
model_path="model/"
paths = get_checkpoint_paths(model_path)
print(paths.config, paths.checkpoint, paths.vocab)
# 下载的数据集的路径
base_dir = 'cnews'
train_dir = os.path.join(base_dir, 'cnews.train.txt')
test_dir = os.path.join(base_dir, 'cnews.test.txt')
val_dir = os.path.join(base_dir, 'cnews.val.txt')
vocab_dir = os.path.join(base_dir, 'cnews.vocab.txt')

数据预处理

def read_file(filename):
    """读取文件数据"""
    contents, labels = [], []
    with open(filename, encoding='utf-8') as f:
        for line in f:
            try:
                label, content = line.strip().split('\t')
                if content:
                    contents.append(content)
                    labels.append(label)
            except:
                pass
    return contents, labels
# 从文件中读取文本内容和对应的分类标签
train_contents, train_labels=read_file(train_dir)
validate_contents, validate_labels=read_file(val_dir)

# 读取预训练模型的“分词器”
token_dict = {}
with codecs.open(paths.vocab, 'r', 'utf8') as reader:
    for line in reader:
        token = line.strip()
        token_dict[token] = len(token_dict)
tokenizer = Tokenizer(token_dict)
tokenizer.tokenize(u'今天天气不错')

注意到“分词器”针对“今天天气不错”会输出：

['[CLS]', '今', '天', '天', '气', '不', '错', '[SEP]']

这里也可以看到BERT的“分词”实际上是将句子中的每一个汉字都拆成了单独的一个字符。这实际上也是前面提到的“基于字符的”神经网络模型。2019年已经有文献证明，针对汉字，基于字符的NLP要比基于分词的NLP训练效果要更好，因此BERT的中文NLP也采用了这种技术。

def get_id_segments(contents):
  ids, segments = [],[]
  for sent in contents:
    id, segment = tokenizer.encode(sent)
    ids.append(id)
    segments.append(segment)
  return ids, segments
train_ids,train_segments = get_id_segments(train_contents)
validate_ids, validate_segments = get_id_segments(validate_contents)
print(train_contents[0], train_ids[0])

输出

马晓旭意外受伤让国奥警惕 无奈大雨格外青睐殷家军记者傅亚雨沈阳报道 来到沈阳，国奥队依然没有摆脱雨水的困扰。7月31日下午6点，国奥队的日常训练再度受到大雨的干扰，无奈之下队员们只慢跑了25分钟就草草收场。31日上午10点，国奥队在奥体中心外场训练的时候，天就是阴沉沉的，气象预报显示当天下午沈阳就有大雨，但幸好队伍上午的训练并没有受到任何干扰。下午6点，当球队抵达训练场时，大雨已经下了几个小时，而且丝毫没有停下来的意思。抱着试一试的态度，球队开始了当天下午的例行训练，25分钟过去了，天气没有任何转好的迹象，为了保护球员们，国奥队决定中止当天的训练，全队立即返回酒店。在雨中训练对足球队来说并不是什么稀罕事，但在奥运会即将开始之前，全队变得“娇贵”了。在沈阳最后一周的训练，国奥队首先要保证现有的球员不再出现意外的伤病情况以免影响正式比赛，因此这一阶段控制训练受伤、控制感冒等疾病的出现被队伍放在了相当重要的位置。而抵达沈阳之后，中后卫冯萧霆就一直没有训练，冯萧霆是7月27日在长春患上了感冒，因此也没有参加29日跟塞尔维亚的热身赛。队伍介绍说，冯萧霆并没有出现发烧症状，但为了安全起见，这两天还是让他静养休息，等感冒彻底好了之后再恢复训练。由于有了冯萧霆这个例子，因此国奥队对雨中训练就显得特别谨慎，主要是担心球员们受凉而引发感冒，造成非战斗减员。而女足队员马晓旭在热身赛中受伤导致无缘奥运的前科，也让在沈阳的国奥队现在格外警惕，“训练中不断嘱咐队员们要注意动作，我们可不能再出这样的事情了。”一位工作人员表示。从长春到沈阳，雨水一路伴随着国奥队，“也邪了，我们走到哪儿雨就下到哪儿，在长春几次训练都被大雨给搅和了，没想到来沈阳又碰到这种事情。”一位国奥球员也对雨水的“青睐”有些不解。 [ 101 7716 3236 3195 2692 1912 1358  839 6375 1744 1952 6356 2664 3187
 1937 1920 7433 3419 1912 7471 4712 3668 2157 1092 6381 5442  987  762
 7433 3755 7345 2845 6887 3341 1168 3755 7345 8024 1744 1952 7339  898
 4197 3766 3300 3030 5564 7433 3717 4638 1737 2817  511  128 3299 8176
 3189  678 1286  127 4157 8024 1744 1952 7339 4638 3189 2382 6378 5298
 1086 2428 1358 1168 1920 7433 4638 2397 2817 8024 3187 1937  722  678
 7339 1447  812 1372 2714 6651  749 8132 1146 7164 2218 5770 5770 3119
 1767  511 8176 3189  677 1286 8108 4157 8024 1744 1952 7339 1762 1952
  860  704 2552 1912 1767 6378 5298 4638 3198  952 8024 1921 2218 3221
 7346 3756]

这里主要是将读取到的文章的文本转换成id，这个id才是BERT模型需要的真正的输入。
同时每一个文章长度都是不一样的，这里简单的画了下各文章经过tokenizer后长度的分布：

lengths = [len(i) for i in train_ids]
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
sns.set(style='darkgrid')
sns.set(font_scale=1.5)
plt.rcParams["figure.figsize"] = (10,5)
lengths = [l for l in lengths]
sns.distplot(lengths, kde=False, rug=False)

其中横轴是文章数，纵轴是样本的个数。这里可以看到有相当的文章的"字“数都超过了1000字。但是BERT模型最长也就只能输入512，同时考虑到Colab上GPU的内存的限制，真正输入的文章长度要更短。
2019年已经有论文对文本分类要截取的“字数”进行了讨论（原文），论文针对IMDB上的影评，考虑了三种情况：

• 截取文章头部510个token
• 截取文章尾部510个token
• 截取文章头部128个token和尾部382个token。

发现第三种情况效果是最好的。
由于我们处理的新闻文章，新闻文章的特点一般都是开宗明义，所以我这里只取了文章头部128个token，因此代码：

MAX_LENGTH = 128
train_ids = pad_sequences(train_ids, maxlen=MAX_LENGTH, dtype="long",
                          value=0, truncating="post", padding="post")
train_segments = pad_sequences(train_segments, maxlen=MAX_LENGTH, dtype="long",
                          value=0, truncating="post", padding="post")
validate_ids = pad_sequences(validate_ids, maxlen=MAX_LENGTH, dtype="long",
                          value=0, truncating="post", padding="post")
validate_segments = pad_sequences(validate_segments, maxlen=MAX_LENGTH, dtype="long",
                          value=0, truncating="post", padding="post")

这里用了keras自带的pad_sequences函数。这个函数对超过指定长度的会截断，没到指定长度的会补0，返回numpy数组。这样就不用自己hard code了。
此外，还要将读取到的文章分类也要转为数字：

distinct_labels = list(set(train_labels + validate_labels))
dict_labels = {label : id for id, label in enumerate(distinct_labels)}
train_labelids, validate_labelids = [dict_labels[label] for label in train_labels],[dict_labels[label] for label in validate_labels]

加载BERT预训练模型

加载BERT预训练模型，实际上调用的就是keras-bert的函数就可以。

bert_model = load_trained_model_from_checkpoint(paths.config.replace('.json','_rbt3.json'), paths.checkpoint, seq_len=None)
for i,layer in enumerate(bert_model.layers):
  layer.trainable = True

调用完后，BERT各层默认是constant，即不参与接下来的训练。从实际使用的情况看，如果只训练自己添加的那几层，几乎达不到分类的效果，因此这里我们采用了所有层都参与训练。训练的层越多，理论上效果越好，但也要考虑到过拟合、以及对硬件资源占用的问题。（之前尝试采用“高配旗舰版”的24个transformer layer的BERT，把gpu搞崩了好几次）由于我们的这个BERT的模型比较精简，因此所有层参与训练的问题不大。

修改BERT模型

修改BERT模型，即前面提到的代码

inputs = bert_model.inputs[:2]
x=bert_model.layers[-1].output
# if returns tuple, then we are using keras lib
# if returns KerasHistory, then we are using tensorflow.python.keras lib
print(type(x._keras_history)) 
x=keras.layers.Lambda(lambda x: x[:, 0], name='slice')(x)
x=keras.layers.Dense(units=3072, activation=keras.backend.tanh)(x)
x=keras.layers.Dropout(rate=0.1,seed=2019)(x)
x=keras.layers.Dense(units=num_labels, activation=keras.backend.softmax)(x)
model=keras.Model(inputs, x)
model.compile(optimizer=keras.optimizers.Adam(1e-4),
			loss=keras.losses.sparse_categorical_crossentropy,
            metrics=[keras.metrics.sparse_categorical_accuracy]
             )
model.summary()

会输出

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
Input-Token (InputLayer)        (None, None)         0                                            
__________________________________________________________________________________________________
Input-Segment (InputLayer)      (None, None)         0                                            
__________________________________________________________________________________________________
Embedding-Token (TokenEmbedding [(None, None, 768),  16226304    Input-Token[0][0]                
__________________________________________________________________________________________________
Embedding-Segment (Embedding)   (None, None, 768)    1536        Input-Segment[0][0]              
__________________________________________________________________________________________________
Embedding-Token-Segment (Add)   (None, None, 768)    0           Embedding-Token[0][0]            
                                                                 Embedding-Segment[0][0]          
__________________________________________________________________________________________________
Embedding-Position (PositionEmb (None, None, 768)    393216      Embedding-Token-Segment[0][0]    
__________________________________________________________________________________________________
Embedding-Dropout (Dropout)     (None, None, 768)    0           Embedding-Position[0][0]         
__________________________________________________________________________________________________
Embedding-Norm (LayerNormalizat (None, None, 768)    1536        Embedding-Dropout[0][0]          
__________________________________________________________________________________________________
Encoder-1-MultiHeadSelfAttentio (None, None, 768)    2362368     Embedding-Norm[0][0]             
__________________________________________________________________________________________________
Encoder-1-MultiHeadSelfAttentio (None, None, 768)    0           Encoder-1-MultiHeadSelfAttention[
__________________________________________________________________________________________________
Encoder-1-MultiHeadSelfAttentio (None, None, 768)    0           Embedding-Norm[0][0]             
                                                                 Encoder-1-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-1-MultiHeadSelfAttentio (None, None, 768)    1536        Encoder-1-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-1-FeedForward (FeedForw (None, None, 768)    4722432     Encoder-1-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-1-FeedForward-Dropout ( (None, None, 768)    0           Encoder-1-FeedForward[0][0]      
__________________________________________________________________________________________________
Encoder-1-FeedForward-Add (Add) (None, None, 768)    0           Encoder-1-MultiHeadSelfAttention-
                                                                 Encoder-1-FeedForward-Dropout[0][
__________________________________________________________________________________________________
Encoder-1-FeedForward-Norm (Lay (None, None, 768)    1536        Encoder-1-FeedForward-Add[0][0]  
__________________________________________________________________________________________________
Encoder-2-MultiHeadSelfAttentio (None, None, 768)    2362368     Encoder-1-FeedForward-Norm[0][0] 
__________________________________________________________________________________________________
Encoder-2-MultiHeadSelfAttentio (None, None, 768)    0           Encoder-2-MultiHeadSelfAttention[
__________________________________________________________________________________________________
Encoder-2-MultiHeadSelfAttentio (None, None, 768)    0           Encoder-1-FeedForward-Norm[0][0] 
                                                                 Encoder-2-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-2-MultiHeadSelfAttentio (None, None, 768)    1536        Encoder-2-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-2-FeedForward (FeedForw (None, None, 768)    4722432     Encoder-2-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-2-FeedForward-Dropout ( (None, None, 768)    0           Encoder-2-FeedForward[0][0]      
__________________________________________________________________________________________________
Encoder-2-FeedForward-Add (Add) (None, None, 768)    0           Encoder-2-MultiHeadSelfAttention-
                                                                 Encoder-2-FeedForward-Dropout[0][
__________________________________________________________________________________________________
Encoder-2-FeedForward-Norm (Lay (None, None, 768)    1536        Encoder-2-FeedForward-Add[0][0]  
__________________________________________________________________________________________________
Encoder-3-MultiHeadSelfAttentio (None, None, 768)    2362368     Encoder-2-FeedForward-Norm[0][0] 
__________________________________________________________________________________________________
Encoder-3-MultiHeadSelfAttentio (None, None, 768)    0           Encoder-3-MultiHeadSelfAttention[
__________________________________________________________________________________________________
Encoder-3-MultiHeadSelfAttentio (None, None, 768)    0           Encoder-2-FeedForward-Norm[0][0] 
                                                                 Encoder-3-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-3-MultiHeadSelfAttentio (None, None, 768)    1536        Encoder-3-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-3-FeedForward (FeedForw (None, None, 768)    4722432     Encoder-3-MultiHeadSelfAttention-
__________________________________________________________________________________________________
Encoder-3-FeedForward-Dropout ( (None, None, 768)    0           Encoder-3-FeedForward[0][0]      
__________________________________________________________________________________________________
Encoder-3-FeedForward-Add (Add) (None, None, 768)    0           Encoder-3-MultiHeadSelfAttention-
                                                                 Encoder-3-FeedForward-Dropout[0][
__________________________________________________________________________________________________
Encoder-3-FeedForward-Norm (Lay (None, None, 768)    1536        Encoder-3-FeedForward-Add[0][0]  
__________________________________________________________________________________________________
slice (Lambda)                  (None, 768)          0           Encoder-3-FeedForward-Norm[0][0] 
__________________________________________________________________________________________________
dense_11 (Dense)                (None, 3072)         2362368     slice[0][0]                      
__________________________________________________________________________________________________
dropout_6 (Dropout)             (None, 3072)         0           dense_11[0][0]                   
__________________________________________________________________________________________________
dense_12 (Dense)                (None, 10)           30730       dropout_6[0][0]                  
==================================================================================================
Total params: 40,279,306
Trainable params: 40,279,306
Non-trainable params: 0

注意最后三行的提示，即所有的参数都参与了训练。
这里还有个tip，构建模型时，如果直接将代码

x=bert_model.layers[-1].output

替换成

x=bert_model(inputs)

这样不会对model实际训练造成影响，但是在打印model.summary()时，bert内部的各层只会显示为一层bert_model。这样也就不方便查看了。

开始训练

batch_size = 64
train_size = (train_ids.shape[0] // batch_size) * batch_size
validate_size = (validate_ids.shape[0] // batch_size) * batch_size

class TrainHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.train_loss = []
        self.train_acc = []

    def on_batch_end(self, batch, logs={}):
        self.train_loss.append(logs.get('loss'))
        self.train_acc.append(logs.get('sparse_categorical_accuracy'))

history = TrainHistory()


model.fit(x=[train_ids[:train_size], train_segments[:train_size]],
        y=np.array(train_labelids[:train_size]),
        validation_data=[[validate_ids[:validate_size], validate_segments[:validate_size]],np.array(validate_labelids[:validate_size])],
        batch_size = batch_size,
        callbacks=[history])

Keras最让人感到激动的一点是，训练只需要一行代码fit就解决了。这里实际上可以直接将训练集和测试集扔进去也没有问题。之所以我把输入截断成batch_size的整数倍，是由于TPU的输入是这样要求的。但由于TPU试验失败，但是代码还是保留了。
另外，为了记录训练时的loss和accuracy，自定义了TransHistory类。在完成训练后，就可以画图表示loss和accuracy的收敛过程了。

import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
d = {'batches': np.array(range(len(history.train_loss))), 
     'loss': np.array(history.train_loss),
     'accuracy': np.array(history.train_acc)}
baseline = [(100, 0.88, 0.7344),
            (200, 0.38,0.9219), 
            (300,0.22, 0.9219),
            (400,0.24,0.9062),
            (500,0.16,0.9688),
            (600,0.084,0.9688),
            (700,0.21,0.9375)]
baseline = np.array(baseline)
df = pd.DataFrame(baseline, columns=("batches","loss","accuracy"))


plt.figure()
sns.set(style='darkgrid')
sns.lineplot(x='batches', y='loss', data=pd.DataFrame(d), label='bert')
sns.lineplot(x='batches', y='loss', data=df,label='cnn')
sns.set(style='darkgrid')
plt.legend()
plt.figure()
sns.lineplot(x='batches', y='accuracy', data=pd.DataFrame(d), label='bert')
sns.lineplot(x='batches', y='accuracy', data=df, label='cnn')
plt.legend()

训练结果输出

Train on 49984 samples, validate on 4992 samples
Epoch 1/1
49984/49984 [==============================] - 199s 4ms/step - loss: 0.1682 - sparse_categorical_accuracy: 0.9495 - val_loss: 0.1362 - val_sparse_categorical_accuracy: 0.9619

这里为了节省时间，只训练了一个epoch，使用google colab的GPU 3分多钟就跑完了，测试集上准确率达到96.19%，而github上采用cnn训练了3个epoch，才最多达到94.12%，效果还是很好的。
同时图像上也做了比较：


从图像上看，模型在100个batch左右，即输入6400个样本的时候，就可以训练到准确率90%左右，这样也说明了BERT在下游训练的优势还是很明显的。这也是最近两年bert如此火的原因吧。