BERT代码实现

前段时间实现了transformer，用李沐老师的话来讲其实bert可以简单理解为缩水版的transformer，transformer有encoder和decoder，bert去掉了decoder，改为用N个encoder堆叠在一起，拿最后一个encoder的输出直接做预训练任务。老规矩，先把大体框架列出来，然后根据框架一个一个去实现。

目录

架构

 数据预处理

NSP

MLM：

BERT

Embedding

Encoder

NSP任务

MLM任务

架构

Bert的架构很简单，包括词向量输入，encoder层，NSP（下一句预测任务）和MLM（掩码词预测任务），如下图

 其中，bert的embedding由三个部分组成，分别是词向量，位置向量，句向量（主要用于NSP任务），所以Embedding的架构如下

bert的encoder和transformer的 encoder一样，由多头注意力和前馈神经网络组成

 数据预处理

首先对数据建立两个词表，分别是词对应序号的word_2_id和序号对应词的id_2_word，BERT使用[CLS]作为一个句子的开头，用[SEP]作为句子的结尾，[MASK]主要为了MLM任务中进行掩码，[PAD]用作填充，所以词表初始化如下

word_2_id = {'PAD': 0, 'CLS': 1, 'SEP': 2, 'MASK':3}
id_2_word = {0: 'PAD', 1: 'CLS', 2: 'SEP', 3:'MASK'}
           

打乱数据的词语，再将数据里的词语添加到词表中

shuffle(sentences)
    for i in range(len(sentences)):
        if sentences[i] not in word_2_id:
            word_2_id[sentences[i]] = len(word_2_id)
            id_2_word[len(id_2_word)] = sentences[i]
    return word_2_id, id_2_word
           

 接下来为NSP和MLM任务要用到的数据作准备

NSP

在数据中随机取出两个句子（注意这两个句子不能是同一个），然后使用[CLS]和[SEP]对两个句子拼接成一个句子对

num_a = randint(0, len(sentences)-1)
    while True:
        num_b = randint(0, len(sentences)-1)
        if num_b != num_a:
            break
    a_and_b = 'CLS ' + ' '.join(sentences[num_a]) + ' SEP ' + ' '.join(sentences[num_b]) + ' SEP'
           

拼接后的句子应该是：CLS 句子1 SEP 句子2 SEP

num_a代表句子1在数据中的位置，num_b代表句子2在数据中的位置，如果num_a + 1 ==num_b，即说明句子1是句子2的上一句，此时将这个句子对设为True标志

if num_a == num_b + 1 and positive != batch_size/2:
        IsNext.append(True)
    elif num_a != num_b and negative != batch_size/2:
        IsNext.append(False)
    else:
        continue
           

拼接好后，句子对中的句子1用0表示，句子2用1表示，这便是句向量

seg_ids = [0] * (len(sentences[num_a]) + 2) + [1] * (len(sentences[num_b]) + 1)
           

MLM：

MLM任务中，对于一个句子对中的单词，有80%的概率将其替换为[MASK]标志，10%的概率使用该句子对的其他单词替换，10%的概率保持不变。这里注意mask_num表示一个句子最多能有多少个词被替换，并且在选择句子中的词和选择要替换的词时，不能是特殊符号[CLS]，[SEP]，和[MASK]

for i in range(mask_num):
    while True:
        num = randint(0,len(token_ids)-1)
        if token_ids[num] != 0 and token_ids[num] != 1 and token_ids[num] != 2 and token_ids[num] != 3:
        break
    mask_token.append(token_ids[num])
    if randint(0,1)<0.8:
        token_ids[num] = 3
     elif randint(0,1)<0.5:
        while True:
            num1 = randint(0, len(token_ids) - 1)
            if token_ids[num1] != 0 and token_ids[num1] != 1 and token_ids[num1] != 2 and token_ids[num] != 3:
                break
        token_ids[num] = token_ids[num1]
    mask_position.append(num)
           

以上是对数据进行预处理，接下来实现BERT的框架。

BERT

根据框架，可以初步列出所需模块

class BERT(nn.Module):
    def __init__(self):
        super(BERT, self).__init__()
        self.embed = Embedding()
        self.encoder = nn.ModuleList(Encoder() for _ in range(n_layers))
        self.cls = CLS()
        self.mlm = MLM()
           

Embedding

BERT的输入词向量包括三个部分，词向量，位置向量，句向量。在数据预处理的时候已经实现了词向量和句向量，只需实现位置向量，位置向量就是句子从头到尾遍历一遍，生成一个[0,1,2...max_seq-1]的张量。

position = np.arange(0, len(input_ids[0]))
position = torch.LongTensor(position)
           

然后将词向量，位置向量，句向量加起来做归一化便得到BERT的输入词向量

input_embed = self.norm(token_embed + position_embed + seg_embed)
           

Encoder

Encoder层包括多头注意力层和前向反馈层，在上一篇博客“Transformer的实现”中讲的很详细了，这里不作过多阐述。

NSP任务

class CLS(nn.Module):
    def __init__(self):
        super(CLS, self).__init__()
        self.linear_a = nn.Linear(Embedding_size, Embedding_size)
        self.tanh = nn.Tanh()
        self.linear_b = nn.Linear(Embedding_size, 2)
    def forward(self,en_output):
        cls_output = self.linear_a(en_output[:, 0])
        cls_output = self.tanh(cls_output)
        cls_output =self.linear_b(cls_output)
        return cls_output
           

经过一个线性层，然后经过tanh激活函数，最后经过一个将原本Embedding_size维度映射为2维的线性层，为什么是2维，因为NSP其实就是一个二分类任务，对下一句是否是下一句，模型只需判断是正确还是错误即可。

MLM任务

class MLM(nn.Module):
    def __init__(self):
        super(MLM, self).__init__()
        self.linear_a = nn.Linear(Embedding_size,Embedding_size)
        self.norm = nn.LayerNorm(Embedding_size)
        self.linear_b = nn.Linear(Embedding_size, len_vocab, bias=False)
        self.softmax = nn.Softmax(dim=2)
    def forward(self,en_output,masked_position):
        masked_position = masked_position.unsqueeze(1).expand(batch_size, Embedding_size, max_mask).transpose(1,2) #[6,5,768]
        mask = torch.gather(en_output, 1, masked_position) #[6,5,768]
        mlm_output = self.linear_a(mask)
        mlm_output = gelu(mlm_output)
        mlm_output = self.norm(mlm_output)
        mlm_output = self.linear_b(mlm_output)
        mlm_output = self.softmax(mlm_output)

        return mlm_output
           

这里需要注意的是BERT使用了gelu作为激活函数，gelu其实是dropout、zoneout、Relus的综合，论文里的计算公式如下

 因为这里做的是预测任务，是多分类任务，所以最后一个线性层将Embedding_size维度映射为词表长度的维度（因为预测词要在词表中选出），然后加上softmax进行分类。但其实去掉softmax模型收敛更快，可能是已经有了gelu作为激活函数了。

以上，BERT框架所需模块已经全部实现了，接下来只需调用即可

def forward(self, input_ids, segment_ids, masked_position):
        embed = self.embed(input_ids,segment_ids) #[6,30,768]
        attn_mask = get_attn_mask(input_ids) #[6,30,30]
        for encoder in self.encoder:
            en_outpput = encoder(embed, attn_mask) #[6,30,768]
        cls_output = self.cls(en_outpput) #[6,2]
        mlm_output = self.mlm(en_outpput, masked_position) #[6,5,29]
        return cls_output, mlm_output
           

BERT的代码和Transformer的很像，基本就是直接将Transformer的encoder和注意力的代码全搬过来。BERT主要是两个预训练任务NSP和MLM，如何处理数据其实就是实现BERT最大的难点了。

完整代码：GitHub - bowspider-man/bert-