技术干货｜百行代码写BERT，昇思MindSpore能力大赏

之前在如何评价华为MindSpore 1.5 ?提到了MindSpore的易用性已经具备百行代码写个BERT的能力，这次补上。BERT作为18年NLP里程碑式的模型，在被无数人追捧的同时也被解构分析了无数次，我力争稍微讲清楚模型的同时，让读者也能Get到MindSpore当前的能力。虽然多少有点广告的嫌疑，但是敬请各位看官听我细细讲来。

 01 

近1000行的BERT实现

动了写这个文章的念头是因为使用MindSpore也写了不少模型，尤其做了些预训练语言模型的复现，其间不断的参考Model Zoo的模型实现，当时就有个疑惑，写个BERT需要将近1000行，MindSpore有这么复杂且难用吗？

官方实现的链接放上（gitee.com/mindspore/models/blob/master/official/nlp/bert/src/bert_model.py），有兴趣的读者可以看看，去掉注释，这份实现仍旧复杂冗长，而且丝毫没有体现出MindSpore宣传的那样——“简单的开发体验”。后来想要去迁移huggingface的checkpoint，自己动手写了一版，发现这冗长的官方实现完全是可以压缩的，且如此复杂的实现会给后来者增加一些困惑，因此就有了百行代码的version。

 02 

BERT模型

BERT是“Bidirectional Encoder Reporesentation from Transfromers”的缩写，也是芝麻街动漫人物的名字（谷歌老彩蛋人了）。

芝麻街中的BERT

标题就已经点出了模型的核心，双向的Transformer Encoder，BERT在GPT和ELMo的基础上，吸纳了二者的优势，并且充分利用Transformer的特征提取能力，在当年狂刷各个评测数据集，成为无法跨越的SOTA模型。接下来我会从最基础的组成部分一点一点的配合公式、图文、代码，用MindSpore完成一个非常轻量的BERT。由于篇幅原因，这里不会对Transformer或预训练语言模型的基础知识进行详细讲解，仅以实现BERT为主线。

 03 

Multi-head Attention

不同于Paper解析，我不会上来就讲BERT和Transformer在Embedding的差异，以及设置的预训练任务，而是从最基础的模块实现开始讲起。首先就是多头注意力(Multi-head Attention)模块。

由于BERT模型的基本骨架完全由Transformer的Encoder构成，所以这里先对Transformer中Self-Attention和Multi-head Attention进行简述。首先是Self-Attention，即论文中的Scaled Dot-product Attention，其公式如下：

这里的Self-Attention由三个输入进行运算，分别是Q(query matrix), K(key matrix), V(value matrix), 其分别由同一个输入经过全连接层进行线性变换而得到。这里的实现可以参照公式进行完全复现，其代码如下：

class ScaledDotProductAttention(Cell):
    def __init__(self, d_k, dropout):
        super().__init__()
        self.scale = Tensor(d_k, mindspore.float32)
        self.matmul = nn.MatMul()
        self.transpose = P.Transpose()
        self.softmax = nn.Softmax(axis=-1)
        self.sqrt = P.Sqrt()
        self.masked_fill = MaskedFill(-1e9)
        if dropout > 0.0:
            self.dropout = nn.Dropout(1-dropout)
        else:
            self.dropout = None

    def construct(self, Q, K, V, attn_mask):
        K = self.transpose(K, (0, 1, 3, 2))
        scores = self.matmul(Q, K) / self.sqrt(self.scale) # scores : [batch_size x n_heads x len_q(=len_k) x len_k(=len_q)]
        scores = self.masked_fill(scores, attn_mask) # Fills elements of self tensor with value where mask is one.
        attn = self.softmax(scores)
        context = self.matmul(attn, V)
        if self.dropout is not None:
            context = self.dropout(context)
        return context, attn

其中，Q*K^T并进行缩放后，做了一步masked_fill的操作，是参考Pytorch版本实现，将计算所得的结果和初始输入序列中Padding为0的对应位置的数值进行替换，替换为接近于0的数，如上述代码中的-1e-9。此外还有增加模型鲁棒性的Dropout操作。

Multi-head Attention

在完成基本的Scaled Dot-product Attention后，再来看多头注意力机制的实现。所谓多头，实际上是将原本单一的Q、K、V投影为h个Q', K', V'。由此在不改变计算量的情况下，增强了模型的泛化能力。这里既可以将其视为多个head在模型内部的集成(ensamble)，也可以等价视作卷积操作中的多通道(channel)，实际上Multi-head Attention也不无借鉴CNN的味道(多年前听刘铁岩老师在讲习班中提及)。下面来看实现部分：

class MultiHeadAttention(Cell):
    def __init__(self, d_model, n_heads, dropout):
        super().__init__()
        self.n_heads = n_heads
        self.W_Q = Dense(d_model, d_model)
        self.W_K = Dense(d_model, d_model)
        self.W_V = Dense(d_model, d_model)
        self.linear = Dense(d_model, d_model)
        self.head_dim = d_model // n_heads
        assert self.head_dim * n_heads == d_model, "embed_dim must be divisible by num_heads"
        self.layer_norm = nn.LayerNorm((d_model, ), epsilon=1e-12)
        self.attention = ScaledDotProductAttention(self.head_dim, dropout)
        # ops
        self.transpose = P.Transpose()
        self.expanddims = P.ExpandDims()
        self.tile = P.Tile()
        
    def construct(self, Q, K, V, attn_mask):
        # q: [batch_size x len_q x d_model], k: [batch_size x len_k x d_model], v: [batch_size x len_k x d_model]
        residual, batch_size = Q, Q.shape[0]
        q_s = self.W_Q(Q).view((batch_size, -1, self.n_heads, self.head_dim)) 
        k_s = self.W_K(K).view((batch_size, -1, self.n_heads, self.head_dim)) 
        v_s = self.W_V(V).view((batch_size, -1, self.n_heads, self.head_dim)) 
        # (B, S, D) -proj-> (B, S, D) -split-> (B, S, H, W) -trans-> (B, H, S, W)
        q_s = self.transpose(q_s, (0, 2, 1, 3)) # q_s: [batch_size x n_heads x len_q x d_k]
        k_s = self.transpose(k_s, (0, 2, 1, 3)) # k_s: [batch_size x n_heads x len_k x d_k]
        v_s = self.transpose(v_s, (0, 2, 1, 3)) # v_s: [batch_size x n_heads x len_k x d_v]

        attn_mask = self.expanddims(attn_mask, 1)
        attn_mask = self.tile(attn_mask, (1, self.n_heads, 1, 1)) # attn_mask : [batch_size x n_heads x len_q x len_k]
        
        # context: [batch_size x n_heads x len_q x d_v], attn: [batch_size x n_heads x len_q(=len_k) x len_k(=len_q)]
        context, attn = self.attention(q_s, k_s, v_s, attn_mask)
        context = self.transpose(context, (0, 2, 1, 3)).view((batch_size, -1, self.n_heads * self.head_dim)) # context: [batch_size x len_q x n_heads * d_v]
        output = self.linear(context) 
        return self.layer_norm(output + residual), attn # output: [batch_size x len_q x d_model]

Q，K，V首先经过全连接层（Dense）进行线性变换，然后经过reshape（view）切换为多头，继而进行相应的转置满足送入ScaledDotProductAttention的需要。最后将获得的输出进行拼接，注意这里并没有显式的进行Concat操作，而是直接通过view，将context的shape[-1]还原为heads*hidden_size的大小。此外，最后return时加入了Add&Norm操作，即Encoder结构中对应的残差和Norm计算。这里不进行详述，参见下一节。

 04 

Transformer Encoder

在完成基础的Multi-head Attention模块后，可以将其余部分完成，构造单层的Encoder。这里先对单层Encoder的结构进行简单说明，Transformer Encoder由Poswise Feed Forward Layer和Multi-head Attention Layer构成，并且每个Layer的输入和输出做了Residual运算(即: y = f(x) + x), 来保证加深神经网络层数不会产生退化问题，以及Layer Norm来满足深层神经网络可训练(缓解梯度消失和梯度爆炸)。这里为何使用Layer Norm而非Batch Norm可自行搜索，也是Transformer模型构造的一个有趣的trick。

Transformer Encoder

讲完Encoder的结构，需要将缺少的Poswise Feed Forward Layer进行实现，同时与Multi-head Attention Layer相仿，将Residual和Layer Norm集成到一起，代码实现如下：

class PoswiseFeedForwardNet(Cell):
    def __init__(self, d_model, d_ff, activation:str='gelu'):
        super().__init__()
        self.fc1 = Dense(d_model, d_ff)
        self.fc2 = Dense(d_ff, d_model)
        self.activation = activation_map.get(activation, nn.GELU())
        self.layer_norm = nn.LayerNorm((d_model,), epsilon=1e-12)

    def construct(self, inputs):
        residual = inputs
        outputs = self.fc1(inputs)
        outputs = self.activation(outputs)
        
        outputs = self.fc2(outputs)
        return self.layer_norm(outputs + residual)

将Multi-head Attention Layer和Poswise Feed Forward Layer连接即可获得Encoder：

class BertEncoderLayer(Cell):
    def __init__(self, d_model, n_heads, d_ff, activation, dropout):
        super().__init__()
        self.enc_self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.pos_ffn = PoswiseFeedForwardNet(d_model, d_ff, activation)

    def construct(self, enc_inputs, enc_self_attn_mask):
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)
        enc_outputs = self.pos_ffn(enc_outputs)
        return enc_outputs, attn

而后根据配置的层数、hidden_size, head数等参数，将n层Encoder依次连接，即可完成BERT的Encoder，这里使用nn.CellList容器进行实现：

class BertEncoder(Cell):
    def __init__(self, config):
        super().__init__()
        self.layers = nn.CellList([BertEncoderLayer(config.hidden_size, config.num_attention_heads, config.intermediate_size, config.hidden_act, config.hidden_dropout_prob) for _ in range(config.num_hidden_layers)])

    def construct(self, inputs, enc_self_attn_mask):
        outputs = inputs
        for layer in self.layers:
            outputs, enc_self_attn = layer(outputs, enc_self_attn_mask)
        return outputs

05 

构造BERT

在完成了Encoder后，可以开始组装完整的BERT模型。前述章节的内容都是Transformer Encoder结构的实现，而BERT模型的核心创新或差异则主要在Transformer backbone以外。首先是对Embedding的处理。

如图所示，文本输入后送入BERT模型的Embedding获得隐层表示由三种不同的Embedding加和而得，其中包括：

• Token Embeddings：即最常见的词向量，其中第一个占位符为[CLS]，用于后续编码后表达整条输入文本的编码，用于分类任务(因此成为CLS，即classifier)。此外还有[SEP]占位符用于分隔同一条输入的两个不同的句子，以及[PAD]表示Padding。

• Segment Embedding：用以区分同一条输入的两个不同句子。该Embedding的加入是为了进行Next Sentence Predict任务。

• Position Embedding：与Transformer一样，其无法像LSTM天然保留了位置信息，则需要手动对位置信息进行编码，这里的区别在于Transformer使用了三角函数，而这里则直接将位置对应的index送入Embedding层获取编码。（二者并无本质区别，且后者更简单直接）

分析完三种不同的Embedding，直接使用nn.Embedding即可完成该部分，对应代码如下：

class BertEmbeddings(Cell):
    def __init__(self, config):
        super().__init__()
        self.tok_embed = Embedding(config.vocab_size, config.hidden_size)
        self.pos_embed = Embedding(config.max_position_embeddings, config.hidden_size)
        self.seg_embed = Embedding(config.type_vocab_size, config.hidden_size)
        self.norm = nn.LayerNorm((config.hidden_size,), epsilon=1e-12)

    def construct(self, x, seg):
        seq_len = x.shape[1]
        pos = mnp.arange(seq_len) # mindspore.numpy
        pos = P.BroadcastTo(x.shape)(P.ExpandDims()(pos, 0))
        seg_embedding = self.seg_embed(seg)
        tok_embedding = self.tok_embed(x)
        embedding = tok_embedding + self.pos_embed(pos) + seg_embedding
        return self.norm(embedding)

这里使用了mindspore.numpy.arange来生成位置index，其余均为简单的调用和矩阵加。

在完成Embedding层后，将Encoder和输出的pooler组合，即可构成完整的BERT模型，其代码如下：

class BertModel(Cell):
    def __init__(self, config):
        super().__init__(config)
        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)
        self.pooler = Dense(config.hidden_size, config.hidden_size, activation='tanh')
        
    def construct(self, input_ids, segment_ids):
        outputs = self.embeddings(input_ids, segment_ids)
        enc_self_attn_mask = get_attn_pad_mask(input_ids, input_ids)
        outputs = self.encoder(outputs, enc_self_attn_mask)
        h_pooled = self.pooler(outputs[:, 0]) 
        return outputs, h_pooled

这里使用一个全连接层，对位置为0的输出进行pooler操作，即对应[CLS]占位符的输入文本表示，以供后续的分类任务使用。

 06 

BERT预训练任务

BERT模型的精髓在于任务设计而非模型结构，是大家对其Paper的共识。BERT共设计了两个预训练任务，来完成无监督条件下的语言模型训练(实际上并非无监督)。

1. Next Sentence Predict

首先先对实现较为简单的NSP任务进行分析。加入NSP任务主要是为了针对QA或NLI等输入句子数为2个的下游任务，增强模型在此类任务的能力。该预训练任务顾名思义，将句子A和B拼接作为输入，其中B有一半为正确情况，是A的下一句，另一半则随机选取非下一句的文本。预测任务则是二分类，预测B是否为A的下一句。具体实现如下：

class BertNextSentencePredict(Cell):
    def __init__(self, config):
        super().__init__()
        self.classifier = Dense(config.hidden_size, 2)

    def construct(self, h_pooled):
        logits_clsf = self.classifier(h_pooled)
        return logits_clsf

2. Masked Language Model

Mask的Token。该任务有别于传统语言模型(或GPT的语言模型)，在于其为双向，即：

以此为目标函数，通过上下文预测被Mask的Token，天然符合完形填空的形态。

这里不涉及数据预处理的内容，就不对Mask和替换比例进行详述了。对应的实现比较简单，实际上是Dense+activation+LayerNorm+Dense，实现如下：

class BertMaskedLanguageModel(Cell):
    def __init__(self, config, tok_embed_table):
        super().__init__()
        self.transform = Dense(config.hidden_size, config.hidden_size)
        self.activation = activation_map.get(config.hidden_act, nn.GELU())
        self.norm = nn.LayerNorm((config.hidden_size, ), epsilon=1e-12)
        self.decoder = Dense(tok_embed_table.shape[1], tok_embed_table.shape[0], weight_init=tok_embed_table)

    def construct(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.norm(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states

将两个Task组合，即可完成预训练的BERT模型：

class BertForPretraining(Cell):
    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)
        self.nsp = BertNextSentencePredict(config)
        self.mlm = BertMaskedLanguageModel(config, self.bert.embeddings.tok_embed.embedding_table)

    def construct(self, input_ids, segment_ids):
        outputs, h_pooled = self.bert(input_ids, segment_ids)
        nsp_logits = self.nsp(h_pooled)
        mlm_logits = self.mlm(outputs)
        return mlm_logits, nsp_logits

行文至此，用MindSpore实现整个BERT模型就完成了，可以看到，每个模块都和公式或图示能够完全对应，且单个模块的实现均在10-20行左右，总体实现代码在150-200行之间，相较于Model Zoo的800+代码，实在简洁。

 07 

前后对比

由于官方实现实在冗长，这里选择一部分代码的截图对比来看

左侧为官方实现的BERTModel，右侧为上述实现的集成，同样的BERT模型可以通过100多行代码进行简洁实现。由此可见，MindSpore在经过多版本迭代后，其本身的算子支持度和前端表达的易用性已经逐步趋向于完善，百行代码实现BERT，以前或许只有Pytorch能做到的，如今MindSpore也可以。

当然，官方实现由于是早期版本一直持续维护，应该没有在版本更迭后考虑使用更简洁的方式去完成，但是这会给用户造成MindSpore很难用，要多写很多代码的错觉。在1.2版本发布后，其能力逐步已经能够支撑与Pytorch持平的代码量完成同量级模型，希望这篇文章，能够成为一个小的样例。

 08 

小结

最后还是总结一下。首先，从我个人的使用体会，MindSpore从0.7的将就可用，到1.0的基本完善，再到1.5的易用性提升，在“简单的开发体验”这个目标上，是有质的飞跃的。而ModelZoo毕竟模型众多，且少有人不断重构优化，造成的误解应该不在少数。所以，就以BERT这个里程碑式的模型为例，让大家有一点直观的体会。

此外，对所有的NLPer多啰嗦几句，BERT也就是100来行代码的一个模型，Transformer结构所见即所得，不要畏惧大模型，自己复现一下，不管是做实验发Paper还是面试回答问题，都要得心应手得多。所以，拿着MindSpore写起来吧！

MindSpore官方资料

官方QQ群 : 486831414

官网：https://www.mindspore.cn/

Gitee : https : //gitee.com/mindspore/mindspore

GitHub : https://github.com/mindspore-ai/mindspore

论坛：https://bbs.huaweicloud.com/forum/forum-1076-1.html