结合原理和代码来理解bert模型

写在前面

本文主要记录在学习https://blog.csdn.net/qq_37236745/article/details/108845470这篇博客中的bert模型代码时，我所理解的和学到的东西。本文会结合bert模型原理对这篇博客中的代码进行逐行逐句细致的解析，一些实在看不懂的地方会略过，一些错误的地方希望大家指正。建议在往下看之前先看一下这两篇文章：Transformer 模型详解，图解BERT模型：从零开始构建BERT，对Encoder结构有一个基本的认识。

---

源代码

首先贴上博客中的源代码，其中有我添加的补充注释，如果有兴趣也可以继续往下看更加详细的结合原理的解析。这段代码实现了bert模型的预训练任务，语料是两个人的几句英文对话，方便关注模型本身。建议粘贴到自己的IDE中，跟着我的解析对照着看。

# -*- coding: utf-8 -*-
"""BERT-Torch

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1LVhb99B-YQJ1bGnaWIX-2bgANy78zAAt
"""

'''
  code by Tae Hwan Jung(Jeff Jung) @graykode, modify by wmathor
  Reference : https://github.com/jadore801120/attention-is-all-you-need-pytorch
         https://github.com/JayParks/transformer, https://github.com/dhlee347/pytorchic-bert
'''
import re
import math
import torch
import numpy as np
from random import *
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data

text = (
    'Hello, how are you? I am Romeo.\n' # R
    'Hello, Romeo My name is Juliet. Nice to meet you.\n' # J
    'Nice meet you too. How are you today?\n' # R
    'Great. My baseball team won the competition.\n' # J
    'Oh Congratulations, Juliet\n' # R
    'Thank you Romeo\n' # J
    'Where are you going today?\n' # R
    'I am going shopping. What about you?\n' # J
    'I am going to visit my grandmother. she is not very well' # R
)
# 将每个句子的标点符号去掉，并形成列表
sentences = re.sub("[.,!?\\-]", '', text.lower()).split('\n') # filter '.', ',', '?', '!'
# 将每个句子拆分成多个单词，形成列表
word_list = list(set(" ".join(sentences).split())) # ['hello', 'how', 'are', 'you',...]
word2idx = {'[PAD]' : 0, '[CLS]' : 1, '[SEP]' : 2, '[MASK]' : 3}
# enumerate函数用于将一个列表中的元素形成索引序列，例如：l=['a','b','c']，enumerate(l)=(0,'a'),(1,'b'),(2,'c')
# 将每个词都设置一个索引，和PAD，CLS，SEP，MASK一起形成一个字典，形式为“单词：索引”
for i, w in enumerate(word_list):
    word2idx[w] = i + 4
# 与word2idx相反，其形式为“索引：单词”
idx2word = {i: w for i, w in enumerate(word2idx)}
# 词表大小
vocab_size = len(word2idx)

# 获取每个句子中每个词语在词表中的的索引，形成列表
token_list = list()
for sentence in sentences:
    arr = [word2idx[s] for s in sentence.split()]
    token_list.append(arr)

# BERT Parameters
maxlen = 30 # 表示每个batch中的所有句子都由30个token组成，不够的补PAD
batch_size = 6
max_pred = 5 # max tokens of prediction，最多需要预测多少个词语
n_layers = 6 # 表示encoder layer的数量
n_heads = 12 # 是指Multi-Head-Attention中self-Attention的个数
d_model = 768 # 表示Token Embedding，Segment Embedding、Position Embedding的维度
d_ff = 768*4 # 4*d_model, FeedForward dimension ，表示Encoder layer中全连接层的维度
d_k = d_v = 64  # dimension of K(=Q), V
n_segments = 2 # 表示Decoder input由几句话组成

# 数据预处理部分，需要mask一句话中15%的token，还需要随机拼接两句话
# sample IsNext and NotNext to be same in small batch size
def make_data():
    batch = []
    positive = negative = 0
    while positive != batch_size/2 or negative != batch_size/2:
        tokens_a_index, tokens_b_index = randrange(len(sentences)), randrange(len(sentences)) # sample random index in sentences
        # 取出这两个句子的单词索引
        tokens_a, tokens_b = token_list[tokens_a_index], token_list[tokens_b_index]
        # 将随机选取的两个句子的单词索引拼接在一起，而且加入CLS和SEP标记，此时input_ids中每个元素表示单词在词表中的索引
        input_ids = [word2idx['[CLS]']] + tokens_a + [word2idx['[SEP]']] + tokens_b + [word2idx['[SEP]']]
        # 组成Segment Embedding
        segment_ids = [0] * (1 + len(tokens_a) + 1) + [1] * (len(tokens_b) + 1)

        # MASK LM
        n_pred = min(max_pred, max(1, int(len(input_ids) * 0.15))) # 15 % of tokens in one sentence，确定要mask的单词数量
        # 此时cand_maked_pos表示在input_ids中有哪几个位置可以被mask，这个位置是指在此列表中的位置，而不是在词汇表中的索引。
        cand_maked_pos = [i for i, token in enumerate(input_ids)
                          if token != word2idx['[CLS]'] and token != word2idx['[SEP]']] # candidate masked position，选择出可以mask的位置的索引，因为像CLS和SEP这些位置不可以被mask
        # shuffle将序列中的元素随机排序，实现随机选取单词来mask
        shuffle(cand_maked_pos)
        masked_tokens, masked_pos = [], []
        for pos in cand_maked_pos[:n_pred]:
            # masked_pos 表示要mask的单词在input_ids中的位置，而不是在词表中的索引
            masked_pos.append(pos)
            # masked_tokens表示要mask的单词在此表中的索引，因为input_ids中存的就是选取的两个句子的单词索引
            masked_tokens.append(input_ids[pos])
            if random() < 0.8:  # 80%
                input_ids[pos] = word2idx['[MASK]'] # make mask，进行mask
            elif random() > 0.9:  # 10%
                index = randint(0, vocab_size - 1) # random index in vocabulary，替换为词表中一个随机的单词
                while index < 4: # can't involve 'CLS', 'SEP', 'PAD'
                  index = randint(0, vocab_size - 1)
                input_ids[pos] = index # replace

        # Zero Paddings，使得一个batch中的句子都是相同长度
        n_pad = maxlen - len(input_ids) # 需要补的0的个数
        input_ids.extend([0] * n_pad) # extend函数会在列表末尾一次性添加另一个序列的多个值
        segment_ids.extend([0] * n_pad)

        # Zero Padding (100% - 15%) tokens，补齐mask的序列，保证一个batch中所有句子mask的数量是一样的
        if max_pred > n_pred:
            n_pad = max_pred - n_pred
            masked_tokens.extend([0] * n_pad)
            masked_pos.extend([0] * n_pad)

        # 正样本，即两个句子是相连的
        if tokens_a_index + 1 == tokens_b_index and positive < batch_size/2:
            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, True]) # IsNext
            positive += 1
        # 负样本，而且要保证正负样本的比例是1:1
        elif tokens_a_index + 1 != tokens_b_index and negative < batch_size/2:
            batch.append([input_ids, segment_ids, masked_tokens, masked_pos, False]) # NotNext
            negative += 1
    return batch
# Proprecessing Finished

batch = make_data()
# 将batch中的数据分开，input_ids, segment_ids, masked_tokens, masked_pos, isNext分别存到一个单独的集合中
input_ids, segment_ids, masked_tokens, masked_pos, isNext = zip(*batch)
input_ids, segment_ids, masked_tokens, masked_pos, isNext = \
    torch.LongTensor(input_ids),  torch.LongTensor(segment_ids), torch.LongTensor(masked_tokens),\
    torch.LongTensor(masked_pos), torch.LongTensor(isNext)

class MyDataSet(Data.Dataset):
  def __init__(self, input_ids, segment_ids, masked_tokens, masked_pos, isNext):
    self.input_ids = input_ids
    self.segment_ids = segment_ids
    self.masked_tokens = masked_tokens
    self.masked_pos = masked_pos
    self.isNext = isNext

  def __len__(self):
    return len(self.input_ids)

  def __getitem__(self, idx):
    return self.input_ids[idx], self.segment_ids[idx], self.masked_tokens[idx], self.masked_pos[idx], self.isNext[idx]

loader = Data.DataLoader(MyDataSet(input_ids, segment_ids, masked_tokens, masked_pos, isNext), batch_size, True)

# 将之前补0的位置mask掉，让其不参与运算
def get_attn_pad_mask(seq_q, seq_k):
    # batch_size就是上述定义的6，seq_len即句子长度30
    batch_size, seq_len = seq_q.size()
    # eq(zero) is PAD token
    # seq_q.data.eq(0)可以找出句子中哪些位置是PAD标记法，返回的数据与seq_q的维度相同。然后用unsqueeze来扩充维度
    pad_attn_mask = seq_q.data.eq(0).unsqueeze(1)  # [batch_size, 1, seq_len]
    # 扩充维度后再将其变形为[batch_size, seq_len, seq_len]的维度
    return pad_attn_mask.expand(batch_size, seq_len, seq_len)  # [batch_size, seq_len, seq_len]

# 激活函数
def gelu(x):
    """
      Implementation of the gelu activation function.
      For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
      0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
      Also see https://arxiv.org/abs/1606.08415
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

class Embedding(nn.Module):
    def __init__(self):
        super(Embedding, self).__init__()
        # Embedding模块主要有两个参数，第一个是单词本中的单词个数，第二个是输出矩阵的大小
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model)  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        # LayerNorm是一个归一化层
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, seg):
        # size函数输出矩阵的某个维度
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long)
        # pos = pos.unsqueeze(0)首先将pos扩充为二维矩阵，然后expand_as将pos扩充为与x维度相同的矩阵
        pos = pos.unsqueeze(0).expand_as(x)  # [seq_len] -> [batch_size, seq_len]
        # 计算最终输入的Embedding，此时的Embedding是的维度为[batch_size,max_len,d_model]
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)

# 计算self-Attention的输出
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores : [batch_size, n_heads, seq_len, seq_len]
        # 在归一化的时候，scores中的0也会有一个值，会影响最终的结果，所以将之前补0的位置替换为一个非常小的负数，不让它影响softmax的结果
        scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one.将scores矩阵中attn_mask上值为true所对应的位置填充为-1e9，也就是那些补0的位置
        # 这里dim设置为-1，表示对某一维度进行归一化
        # self-Attention的输出矩阵，是要对一个单词对其他所有单词的attention系数进行归一化，所以是在同一维度上的。不是同一位置(0)，也不是同一列(1)，也不是同一行(2)，所以dim设置为-1
        attn = nn.Softmax(dim=-1)(scores)
        # context的维度为：[batch_size,n_head,seq_len,d_k]
        context = torch.matmul(attn, V)
        return context

class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        # 线性变换矩阵
        self.W_Q = nn.Linear(d_model, d_k * n_heads)
        self.W_K = nn.Linear(d_model, d_k * n_heads)
        self.W_V = nn.Linear(d_model, d_v * n_heads)
    def forward(self, Q, K, V, attn_mask):
        # seq_len表示句子的长度，即q，k，v的行数是句子中的单词书，列数是我们自己设置d_model
        # q: [batch_size, seq_len, d_model], k: [batch_size, seq_len, d_model], v: [batch_size, seq_len, d_model]
        residual, batch_size = Q, Q.size(0)
        # (B, S, D) -proj-> (B, S, D) -split-> (B, S, H, W) -trans-> (B, H, S, W)
        # 多个self-Attention就要生成多维的q,k,v。transpose用于交换矩阵的两个维度
        # view函数中的参数-1表示该维度的维数由其他维度来估算得到，也就是说这个维度的维数会由程序自动计算
        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # q_s: [batch_size, n_heads, seq_len, d_k]
        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # k_s: [batch_size, n_heads, seq_len, d_k]
        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1,2)  # v_s: [batch_size, n_heads, seq_len, d_v]

        # unsqueeze对数据维度进行扩充
        # repeat指在某个维度进行重复，repeat(1,n_heads,1,1)表示在第二个维度上重复n_heads次。使pad_mask保持与q，k，v相同的维度
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # attn_mask : [batch_size, n_heads, seq_len, seq_len]

        # context: [batch_size, n_heads, seq_len, d_v], attn: [batch_size, n_heads, seq_len, seq_len]
        # 计算Multi-Head-Attention的输出，但此时还不是最终输出，还没有将多个self-Attention的输出拼接起来
        context = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)

        # 交换第一维度和第二维度的维数，transpose函数在交换时并不会重新开辟一块内存来存储转换后的数据，而是保持原有数据存放位置不变修改一些行列的对应关系
        # 也就是说，经过transpose后，两个矩阵实际共享同一块内存，修改一个矩阵，另一个矩阵的值也会随之变化。
        # contiguous函数会将转换后的矩阵，按照其维度来从头开辟一块内存，并原模原样存放该矩阵，不再共享内存
        # n_heads * d_v就表示将多个self-Attention的输出矩阵拼接起来
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v) # context: [batch_size, seq_len, n_heads, d_v]
        # 进行线性变换
        output = nn.Linear(n_heads * d_v, d_model)(context)
        # 经归一化后生成最终输出，且输出矩阵与输入矩阵的维度是一样的
        # output+ residual表示残差连接
        return nn.LayerNorm(d_model)(output + residual) # output: [batch_size, seq_len, d_model]

# feedforward是一个两层的全连接层，第一层使用gelu激活函数，第二层不使用激活函数
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        # (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_ff) -> (batch_size, seq_len, d_model)
        return self.fc2(gelu(self.fc1(x)))

class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    # 每层Encoder都要先经过Multi-Head Attention，再经过feed-forword。而且两者上方都有一个Norm层，用来对每层的激活值进行归一化
    def forward(self, enc_inputs, enc_self_attn_mask):
        enc_outputs = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,V
        enc_outputs = self.pos_ffn(enc_outputs) # enc_outputs: [batch_size, seq_len, d_model]
        return enc_outputs

class BERT(nn.Module):
    def __init__(self):
        super(BERT, self).__init__()
        self.embedding = Embedding()
        # 建立6层的Encoder
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])
        # Sequential是一个有序的容器，神经网络的各种模块在这里面被顺序添加执行
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.Dropout(0.5),
            nn.Tanh(),
        )
        self.classifier = nn.Linear(d_model, 2)
        self.linear = nn.Linear(d_model, d_model)
        self.activ2 = gelu
        # fc2 is shared with embedding layer
        embed_weight = self.embedding.tok_embed.weight
        self.fc2 = nn.Linear(d_model, vocab_size, bias=False)
        self.fc2.weight = embed_weight

    def forward(self, input_ids, segment_ids, masked_pos):
        # 得到transformer的输入Embedding
        output = self.embedding(input_ids, segment_ids) # [bach_size, seq_len, d_model]
        enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids) # [batch_size, maxlen, maxlen]
        for layer in self.layers:
            # output: [batch_size, max_len, d_model]
            # 前一层Encoder的输出作为后一层Encoder的输入
            output = layer(output, enc_self_attn_mask)
        # it will be decided by first token(CLS)
        # output[: ,0]取出每个句子中CLS的所有attention系数
        h_pooled = self.fc(output[:, 0]) # [batch_size, d_model]
        # 用第一个位置CLS的output丢进Linear classifier来预测一个class
        logits_clsf = self.classifier(h_pooled) # [batch_size, 2] predict isNext

        # masked_pos是每句话中要预测的单词的位置，将去扩充到与output相同的维度
        masked_pos = masked_pos[:, :, None].expand(-1, -1, d_model) # [batch_size, max_pred, d_model]
        # 按照masked_pos的值，抽取出output中对应索引的数据
        h_masked = torch.gather(output, 1, masked_pos) # masking position [batch_size, max_pred, d_model]
        h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, d_model]
        logits_lm = self.fc2(h_masked) # [batch_size, max_pred, vocab_size]
        return logits_lm, logits_clsf

model = BERT()
# 交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器，将bert的模型参数传入进行优化
optimizer = optim.Adadelta(model.parameters(), lr=0.001)
for epoch in range(50):
    for input_ids, segment_ids, masked_tokens, masked_pos, isNext in loader:
        logits_lm, logits_clsf = model(input_ids, segment_ids, masked_pos)
        loss_lm = criterion(logits_lm.view(-1, vocab_size), masked_tokens.view(-1)) # for masked LM
        # mean表示求平均值
        loss_lm = (loss_lm.float()).mean()
        loss_clsf = criterion(logits_clsf, isNext) # for sentence classification
        loss = loss_lm + loss_clsf
        if (epoch + 1) % 5 == 0:
            print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))
        # 清空梯度
        optimizer.zero_grad()
        # 计算反向传播
        loss.backward()
        optimizer.step()

# Predict mask tokens ans isNext
input_ids, segment_ids, masked_tokens, masked_pos, isNext = batch[1]
print(text)
print('================================')
print([idx2word[w] for w in input_ids if idx2word[w] != '[PAD]'])

logits_lm, logits_clsf = model(torch.LongTensor([input_ids]),
                 torch.LongTensor([segment_ids]), torch.LongTensor([masked_pos]))
logits_lm = logits_lm.data.max(2)[1][0].data.numpy()
print('masked tokens list : ',[pos for pos in masked_tokens if pos != 0])
print('predict masked tokens list : ',[pos for pos in logits_lm if pos != 0])

logits_clsf = logits_clsf.data.max(1)[1].data.numpy()[0]
print('isNext : ', True if isNext else False)
print('predict isNext : ',True if logits_clsf else False)

---

代码解析

Bert模型是以Transformer模型中Encoder部分为基础的，我在刚开始阅读代码的过程中，一直对照着这两篇文章（Transformer 模型详解，图解BERT模型：从零开始构建BERT）中给出的Transformer结构解析进行代码的梳理，以下的分析也会借助这篇博客中的内容。

首先从程序的运行流程上对这段代码中关键的点进行解析，解析的顺序可能与代码的编写顺序不一样：

（1）bert模型中重要的参数解释

# BERT Parameters
maxlen = 30 # 表示每个batch中的所有句子都由30个token组成，不够的补PAD
batch_size = 6
max_pred = 5 # max tokens of prediction，最多需要预测多少个词语
n_layers = 6 # 表示encoder layer的数量
n_heads = 12 # 是指Multi-Head-Attention中self-Attention的个数
d_model = 768 # 表示Token Embedding，Segment Embedding、Position Embedding的维度
d_ff = 768*4 # 4*d_model, FeedForward dimension ，表示Encoder layer中全连接层的维度
d_k = d_v = 64  # dimension of K(=Q), V
n_segments = 2 # 表示Decoder input由几句话组成

maxlen：bert的预训练任务包含两部分，一是对随机替换掉的单词进行预测（Masked LM），二是对两句话是否为前后相联的句子进行判断（NextSentence Prediction）。所以bert会将随机的两句话进行拼接，再随机替换掉其中一些单词，而这样形成的单词列表的长度可能不一致，给计算带来麻烦。所以我们要将每个单词列表进行补齐，补齐后的长度统一为maxlen，这样在计算时矩阵的维度就统一了。

batch_size：训练模型时，每次输入的句子个数。

max_pred：预训练任务中有一项为对随机替换掉的单词进行预测，而max_pred为做多要预测多少个单词，也就是说最多要替换掉多少个单词。

n_layers：表示Encoder层的个数，Encoder层在下面进行详解。

n_heads：Encoder层中包含一个Mutil-Head-Attention，这个东西由n_heads个self-Attention组成，self-Attention在下面会有解释。

d_model：Bert模型的输入为Embedding，它是由单词Embedding、位置Embedding、句子Embedding组成，d_model为他们的维度，后面会分别对着三个Embedding进行解释。

d_ff：在Encoder中，包含有一个全连接层，d_ff表示他的维度。

d_k，d_v：在self-Attention中，有三个重要的矩阵Q、K、V，他们其中一个维度为d_k。

n_segments：表示一个input由几句话组成，预训练任务中有一项为对两句话是否为前后相联的句子进行判断，所以这里n_segments为2

（2）make_data函数

这个函数用来对语料进行预处理，生成bert所需要的语料格式，引入了CLS，SEP，MASK，PAD这四个标记。还是要提到bert预训练的任务，一是对随机替换掉的单词进行预测，二是对两句话是否为前后相联的句子进行判断。所以在准备语料数据时，我们要随机替换掉一些单词，并且将随机的两句话拼接起来。这段函数中没有很复杂的模型结构，也没有很复杂的编程技巧，看了注释应该还是容易理解的，但其中有一句比较容易理解错：

cand_maked_pos = [i for i, token in enumerate(input_ids)
                          if token != word2idx['[CLS]'] and token != word2idx['[SEP]']]

enumerate函数用于将一个可迭代的对象，组合为一个索引序列，也就是将input_ids中的每个元素按照顺序添加一个索引。也就是说，i 表示元素在input_ids中的位置索引，token表示该元素。而且，token表示的不是具体的单词，而是单词在此表中的索引（晕了的同学回顾一下上面的内容）。if token != word2idx['[CLS]'] and token != word2idx['[SEP]'] 这句用来限定要mask的单词不可以是CLS和SEP标记。token表示的是索引，所以要和wordidx['CLS']做比较，而不是直接跟“CLS”做比较。最后，cand_maked_pos保存的是在input_ids列表中哪几个位置可以被替换掉，也就是说，cand_maked_pos中的元素是inputs_ids列表的索引，而不是词表的索引。

最后总结一下函数返回的batch中一条数据的组成形式：

input_ids：将两个随机的句子拼接成一个句子，加入CLS和SEP标记，保存其中每个单词在词表中的索引。

segment_ids：指示每个单词属于哪个句子

masked_tokens：保存了被换走的单词在词表中的索引。

masked_tokens：保存了被换走的单词在input_ids中的索引。

最后一个布尔变量表示input_ids中两个句子是否相连。

（3）get_attn_pad_mask函数：

这个函数主要是将input_ids中补0的位置掩盖掉，让他们不参与后面的运算，加快运算速度。

def get_attn_pad_mask(seq_q):
    batch_size, seq_len = seq_q.size()
    print(seq_q.data.eq(0).shape)
    # 输出（6,30）
    print(seq_q.data.eq(0).unsqueeze(1).shape)
    # 输出（6,1,30）
    pad_attn_mask = seq_q.data.eq(0).unsqueeze(1)  # [batch_size, 1, seq_len]
    return pad_attn_mask.expand(batch_size, seq_len, seq_len)  # [batch_size, seq_len, seq_len]

其中eq函数用来判断一个对象中哪些元素为0；unsqueeze函数用来扩充维度，原来seq_q.data.eq(0)的维度是（6,30），经过unsqueeze(1)之后的维度变成了（6,1,30）；expand是用来延长矩阵的，原来维度为（6,1,30）的矩阵变成了（6,30,30）。为什么他的维度要设置为（batch_size，seq_len，seq_len）在下面会有解释。

（4）Embedding类：

这个类用来生成bert模型的输入Embedding，其中包含两个函数，初始化函数__init__()和forword()函数。__init__函数用来生成单词Embedding、位置Embedding、句子Embedding；forward函数相当于__call__函数，可以通过类名来直接调用该函数。

    def __init__(self):
        super(Embedding, self).__init__()
        # Embedding模块主要有两个参数，第一个是单词本中的单词个数，第二个是输出矩阵的大小
        self.tok_embed = nn.Embedding(vocab_size, d_model)  # token embedding
        self.pos_embed = nn.Embedding(maxlen, d_model)  # position embedding
        self.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embedding
        # LayerNorm是一个归一化层
        self.norm = nn.LayerNorm(d_model)

在初始化函数中，使用到了pytorch中的Embedding模型，模型中有两个重要的参数，第一个要指定语料中的单词个数vocab_size，第二个指定输出的词向量的维度d_model。这个Embedding模型的输出的是一个（vocab_size，d_model）维度的词向量矩阵。而LayerNorm是一个归一化层，也就是将数值映射到[0,1]范围内，加快收敛速度。

    def forward(self, x, seg):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long)
        pos = pos.unsqueeze(0).expand_as(x)  # [seq_len] -> [batch_size, seq_len]
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(seg)
        return self.norm(embedding)

x在这里接受的是input_ids，seg接收的是segment_ids。seq_len在这里等于30，arange函数会生成一个从0到29的矩阵；然后unsqueeze将其维度变为（0,30）；然后expand_as会复制pos中的内容，的维度扩充到与x相同。将input_ids、pos、segment_ids放入Embedding模型中生成对应的Embedding，最后相加再归一化即生成了要输入到bert中的Embedding。

（5）ScaledDotProductAttention类；

bert模型是以transformer模型为基础的，而transformer模型又是以self-Attention机制为基础的，首先介绍Attention机制。Attention是用来增强目标字与上下文之间的语义联系，将目标字的原始词向量变成增强语义后的词向量，例如单看“鹄”这个目标字，我们很可能连怎么读都不知道，但有了上下文“鸿鹄之志”，就可以获取更多目标字的语义信息。而且，上下文的不同位置对目标字的语义增强效果也是不用的，例如“鸿鹄”肯定比“鹄之”更容易让我们理解语义，这个就是attention机制的思想。

Attention机制主要用到了三个矩阵Query，Key，Value，都是由输入Embedding线性变换得到的。在Attention中，每个词都有他们原始的Value，attention机制将目标字作为Query、其上下文的各个字作为Key。attention机制会计算Query与Key的相似性，然后把计算结果作为权重，加权融合目标字的Value和其他上下文的Value，形成最后的输出。其计算过程如下图所示：

对于输入文本，我们要计算每一个字的语义增强词向量，因此就要将每个字作为Query，融合语料中所有子的Value，得到各个字的attention输出。在这种情况下Q，K，V都是来自于同一语料文本，因此也叫做self-Attention机制，其示意图如下图所示：

回到代码中，这个类要做的就是接收Q，K，V矩阵，并且计算最后的self-Attention输出，其数学计算公式如下图所示：

    def forward(self, Q, K, V, attn_mask):
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores : [batch_size, n_heads, seq_len, seq_len]
        scores.masked_fill_(attn_mask, -1e9)
        attn = nn.Softmax(dim=-1)(scores)
        context = torch.matmul(attn, V)
        return context

matmul函数用来计算两个tensor的乘积，transpose函数用来交换矩阵的两个维度，用来实现K矩阵的转置。masked_fill_函数会将scores矩阵中attn_mask上值为true所对应的位置填充为-1e9，也就是将原先那些补0的位置设置为-1e9。因为在Softmax函数的公式中有一项为，如果传入0的话会输出1，这样会使那些原先我们补0的位置影响最后的计算结果，所以我们将他们设置为一个非常小的常数，消除误差。scores矩阵的维度为（batch_size，n_heads，seq_len，seq_len），在这里（seq_len，seq_len）这两个维度就可以理解为每个词关于其他所有词的融合attention，这样也就解释了上面提到的为什么attn_mask矩阵的维度为（batch_size，seq_len，seq_len）。最后返回的context矩阵的维度为（batch_size，n_heads，seq_len，d_k）。n_heads维度下面再讲。

（6）Multi-Head-Attention类

上面了解了self-Attention的实现过程，为了增加self-Attention的语义多样性，将不同的self-Attention组合在一起就形成了Multi-Head-Attention。例如“南京市长江大桥”，在一种语义场景下，目标字“长”与“市”组合，形成语义“南京市长(zhǎng)”；在另一种语义场景下，目标字“长”与“江”组合，形成语义“长江大桥”。Multi-Head-Attention机制充分考虑了多种语义场景下目标字的语义融合方式，最后形成了从原始输入Embedding到全文增强语义表示的词向量的转换，这也就解释了上面返回的context矩阵中为什么会有一个n_heads维度，就是为了融合了多个self-Attention。如下图所示：

回到代码中，这个类定义了两个函数__init__和forward，初始化函数中定义了Q，K，V矩阵的线性变换模型。forward函数用来计算输出，其中需要注意Q，K，V矩阵的维度变化情况，以Q矩阵为例：

q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)

首先Q矩阵就是原始的词向量，其维度为（batch_size，seq_len，d_model）；经过self.W_Q()线性变换后，其维度不变；view函数只改变数据的排列方式，view函数中参数为-1的那个维度值是由其他维度计算得到的，在这里因为n_heads与d_k相乘就等于d_model，所以计算得到-1所代表的维度仍为seq_len，经此view函数后矩阵维度变换为（batch_size，seq_len，n_heads，d_k）；transpose函数可以交换矩阵的维度，经此函数后矩阵维度变换为（batch_size，n_heads，seq_len，d_k）。

attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)

attn_mask矩阵为了和上面的scores矩阵进行匹配，也要扩充维度。

context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)

将Q，K，V矩阵传入到上面的Scale的DotProductAttention类中得到self-Attention的输出，这句代码将n_heads个self-Attention的输出融合在了一起。这句代码涉及不少细节，首先transpose函数的作用是交换矩阵的两个维度，但其交换方式不是重新开辟一块内存来存放变换维度后的矩阵，而是保持原来数据在内存中的存放位置不变，通过修改行、列的对应关系来实现交换维度，变换前后的矩阵共享同一块内存。也就是说，经transpose函数变换后的矩阵，与从头开始原模原样定义的相同维度的矩阵是不一样的。而contiguous函数的作用就是，将经过transpose转换后的矩阵，按照其维度重新开辟一块内存来存放数据，不再共享内存。view函数用来从n_heads这个维度将attention的输出数据融合在一起。

output = nn.Linear(n_heads * d_v, d_model)(context)
return nn.LayerNorm(d_model)(output + residual)

attention的数据融合在一起后，经过一个线性变换形成最后的输出，此时的输出与输入的Q，K，V矩阵的维度相同。最后经过残差连接和归一化后返回。

（7）PoswiseFeedForwardNet类：

feedforward层是Encoder的一个组成部分，对Multi-Head-Attention层输出的语义增强后的词向量进行两次线性变换，增强整个模型的表达能力；也可以认为是一个两层的全连接层，第一层使用激活函数，第二层不使用激活函数，对应的公式如下：

（8）Encoder类：

这个类用来连接Multi-Head-Attention层和FeedForward层，bert模型中第一个Encoder接收的输入是原始输入Embedding，后面的Encoder接收的输入是前一个Encoder的输出，这点可以在下面的BERT类中得到体现。回顾上面的内容，得到Encoder的整体结构如下图所示：

（9）BERT类：

将多个Encoder连接起来，就构成了bert模型，模型最终生成的就是句子中每个词经语义增强后的词向量。模型的结构大致如下图所示：

这个类包含__init__()和forward()两个函数，__init__函数用来初始化模型中的组件，forward函数相当于python中的__call__函数，可以通过类名来直接调用。forward通过多个Encoder层得到了最后的输出的词向量，然后使用词向量来完成两个预训练任务，output[:0]表示取出每个句子的“CLS”标记的词向量，因为经过多层的语义增强，CLS标记的词向量中已经包含了两个句子的全部信息，使用CLS可以更加准确的预测两个句子是否前后联系。forward函数返回的是两个任务的loss指标，用于后面的训练，loss的计算过程还不是很懂，这里就不乱解释了。

---

结果分析

此次模型训练迭代了50次，由于语料只有两个人的几句对话，随意分析模型的表现也意义不大，但模型确实完成了两项预训练任务，如下图所示，预测到了被替换的词和两句话是否关联。这段代码可以很好地对语料文本进行建模，只需要将代码中的对话文本替换为你的分词后文本即可。

---

GPU加速

如果想要使用自己的数据集，那么bert训练的时间可能就会很长了，这时候可以使用GPU来加速训练，我在写GPU版本的时候遇到了不少bug，调试bug的过程及最后的代码在这篇里：BERT模型迁移到GPU上的调试经历（pytorch），如果有现成的数据集就可以直接拿来运行了。