keras embedding层_BERT的tf.keras实现

BERT出来也很久了，之前一直都是远远观望，现在因为项目需要，想在BERT的基础上尝试，因此认真研究了一下，也参考了几个BERT的实现代码，最终有了这个文章。

本文适合对Transformer比较了解的同学，如果不太了解Transformer及其相关的知识点，建议预先了解一下。

这里有一些不错的关于Transformers的资料：

• 论文 - Attention is All You Need
• The Illustrated Transformer
• Transformer model for language understanding

当然，BERT也有一些非常棒的资料：

• 论文 - BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
• The Illustrated BERT, ELMo, and co.
• A Visual Guide to Using BERT for the First Time

对于BERT的代码实现，我主要是看了以下几个，大家也可以参考：

• huggingface/transformers
• google-research/bert
• tensorflow/models/official/nlp

第一个huggingface/transformers的实现有pytorch和tensorflow实现，整体上还是非常棒的。主要有以下优点：

• 各个SOTA模型都有实现，并且更新速度很快
• 这个库可以很方便地加载预训练好的模型，然后直接使用
• 可以使用tensorflow训练模型和保存，然后使用pytorch加载，进行后面的运算

当然，我没有选择直接使用它而是自己实现，主要原因是：

• 因为网络原因，模型下载太困难 ，导致一些简单的例子和实验都跑不通
• 我只需要BERT，而不需要其他部分，不必杀鸡用牛刀

后面两个实现google-research/bert和tensorflow/models/official/nlp都是Google自己的实现。代码上只能说是一言难尽。

• google-research/bert使用的是tf.estimator API实现，代码组织还可以，但是还是不如tf.keras API来得清晰，并且有不少是处理TPU上训练的代码，对于我们普通人来说，TPU和我们有啥关系呢？
• tensorflow/models/official/nlp使用的是tf.keras API实现，但是代码实在是弯弯绕绕，活生生写成了tf 1.x的样子，不知道是不是阿三哥提交的代码。。。

于是，有了自己实现的一份代码：

• luozhouyang/transformers-keras

好了，废话不多说。

BERT模型的整体结构

首先，我们掌握一下大致的结构，然后逐步实现每一块。

BERT在结构上比Transformer要简单，因为BERT只用了Transformer的Encoder部分。

懒得找图片了，用文字概括一下模型的要点吧：

• Embedding层，包括token embedding、position embedding和segment embedding，所以喂入网络的输入也不仅仅是tokens这么简单。
• Encoder层，实际上是多层encoder的堆叠。基本上和Transformer里的encoder保持一致。其中的multi-head attention也是一样的。
• 对于预训练模型，有两个任务，一个是Masked Language Model，一个是Next Sentence Prediction二分类任务。

到这里也就能理解为什么BERT叫做预训练语言模型了，因为训练这个模型，使用了以上两个任务，但是我们在BERT的基础上进行微调的时候，是可以丢掉上面两个任务，转而在Encoder层后面接上自己的任务的。

可以发现，BERT至少在模型上，是比较简单的，相对于Transformer来说，至少没有Decoder层啊。

接下来，咱们逐步实现以上各个部分。

首先，安装依赖：

!pip install -q tensorflow==2.0.1

import tensorflow as tf

模型参数

为了让BERT模型参数配置更方便，我们单独把它的模型设置参数独立出一个类。主要配置如下：

class BertConfig(object):

def __init__(self, **kwargs):
        super().__init__()
        self.vocab_size = kwargs.pop('vocab_size', 21128) # vocab size of pretrained model `bert-base-chinese`
        self.type_vocab_size = kwargs.pop('type_vocab_size', 2)
        self.hidden_size = kwargs.pop('hidden_size', 768)
        self.num_hidden_layers = kwargs.pop('num_hidden_layers', 12)
        self.num_attention_heads = kwargs.pop('num_attention_heads', 12)
        self.intermediate_size = kwargs.pop('intermediate_size', 3072)
        self.hidden_activation = kwargs.pop('hidden_activation', 'gelu')
        self.hidden_dropout_rate = kwargs.pop('hidden_dropout_rate', 0.1)
        self.attention_dropout_rate = kwargs.pop('attention_dropout_rate', 0.1)
        self.max_position_embeddings = kwargs.pop('max_position_embeddings', 512)
        self.max_sequence_length = kwargs.pop('max_sequence_length', 512)

应该看名字就知道是啥，这里也就不多解释了 。

Embedding层

刚刚上面也说了，不同于Transformer，它只有token embedding和positional embedding，但是BERT的embedding包括三个：

• token embedding，和transformer是一致的
• position embedding，和transformer不一样，transformer使用三角函数来实现位置编码,而BERT使用绝对位置来进行位置编码
• segment embedding，主要是为了区分不同的序列

BERT预训练模型的输入，是一个序列。但是有一个NSP任务，它是对两个序列的分类任务，通常来说需要两个序列单独输入。那么怎么解决这个矛盾呢？

很简单，把两个序列拼接起来！

这里的拼接，不是直接第二个序列跟在第一个序列后面，也是有技巧的。具体做法为：

• 在序列的开头出加上一个[CLS]标记。这个标记有其他用处，后文会说明。
• 在第一个序列结尾处，增加[SEP]标记
• 在第二个序列结尾出，增加[SEP]标记

但是这样还不够，虽然合成了一个序列，但是对于模型来说，它无法区分哪些token是第一个序列的，哪些是第二个序列的啊。所以，额外引入了一个segment embedding，也就是说，用不同的数字来标记不同序列的token。第一个序列的token标记0，第二个序列的token标记1，以此类推。。。

给个例子直观感受下模型的输入是啥样子：

                    [CLS] bert is awesome . [SEP] i love it [SEP]
   token ids:       100   34   3  6       5 101   2 9    4  101
position ids:       0     1    2  3       4 5     6 7    8  9
 segment ids:       0     0    0  0       0 0     1 1    1  1

应该很清楚了，对于一个序列，三个embedding的输入就是这么简单。都转化为ID之后，在各个的embedding矩阵，直接索引即可获得对应的表示。这个和传统的token embedding是一摸一样的。

position