之前在如何评价华为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