pytorch 实现 bert,附带详细的注释和 transformers 国内下载链接

简介

Bert 是 NLP 领域(甚至是在 DL 领域)最近几年最重要的论文了,其将预训练任务、 attention 发扬光大,开辟了一个非常有趣的研究放方向,甚至后续的很多 cv 网络中(如 vitvilbertmae)都可以看到它的身影。

使用纯 pytorch 实现(无 transformers 等多余依赖): backbone_bert

代码实现

bert 的结构并不复杂,但对于刚入门的同学来说,理解起来还是有一点点麻烦的,我们先拿出 transormer 的结构图来。

bert 只使用了 transformer 的 encoder 部分,也就是下边这一部分。

1、Bert Embedding

对照上边的图,我们先实现第一部分,也就是 input_embeddingpostional_embedding 的部分。

  • input_embeddingsegment_embedding 是随机初始化得到的;
  • postinal_embedding 可以通过初始化得到,也可以通过 sin_cos 的方式得到,效果差不多;
  • transformer 中,segment_id 也被称作 type_idinput_id 也被称作 token_id,都一回事;
  • 代码实现参考 bert_layer.py#L17-L64了;

大家可能会看到这里的 LayerNorm 比较特别,是自己实现的 layer_norm 代码,这块其实结果和 torch.nn.LayerNorm 是没有区别的。但是在效率上,torch.nn.LayerNorm 速度更快一些,可能是 torch 自己做了一个额外的优化导致。

题外话,需要注意 LayerNorm 和 BatchNorm 的区别,面试的时候我经常问 😂 。 LayerNorm 是对每一条数据进行 Norm,而不是每一批数据,这两个很像,但作用纬度不一样。在 NLP 任务中,我们使用 LayerNorm 比较多,因为是:

  • 文本自身是变长的,max_length 为 512 的话,可能大部分的数据都只有几十个字。那么让这几十个字以及大批的 padding 进行 norm 是不合理的。
  • batchNorm 中的 平均值 和 方差,是在训练任务中学到的。 然后推理的时候,根据训练任务中学到的平均值和方法来使用,比如 cv 中常见的 transforms.Normalize。如果使用 LayerNorm 的话,就不需要提前计算好平均值和方法,每句话输入进来的时候,单独计算就可以了。对于变长文本预测来说,这样其实更合理一些。
  • 自己实现 layerNorm 还可以方便后续进行一些细小的优化。可参考 https://iii.run/archives/fae41911210f.html

2、Multi-Head Attention

接下来,我们实现第二个部分 Multi-Head Attention 多头注意力机制,我们先看单纯的 点积Attention 结构。

这一部分的代码比较长,可以直接参考 bert_layer.py#L67-L190,基本上都有注释。 我们知道,多头注意力中每个头可以注意到不同的内容,需要实现一个高效的多头机制。而对纬度直接进行调整,从而得到多个头的方式非常高效。

也就是这里的实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def transpose_for_scores(self, x):
"""
这个函数的名字起的比较让人费解
举个例子,以标准的 bert-base 的 query 来说, 输入的 x 纬度为 [batch_size, query_len, hidden_size]
hidden_size 为 768
num_attention_heads 为 12
attention_head_size 为 768 / 12 = 64
new_x_shape = [batch_size, query_len] + [12, 64] 即 [batch_size, query_len, num_attention_heads, attention_head_size]
换句话来说,这个函数其实是把每个 token 的向量都分成了 12 份,给每个注意力头准备了 64d 的数。
"""

new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
x = x.view(*new_x_shape)
return x.permute(0, 2, 1, 3)

在 q*k 的时候,num_attention_heads 应该是不感知的,所以需要将 num_attention_heads 调整到第二个纬度上来。

自此就实现了 Scaled Dot-Product Attention 的部分。

3、Add & Norm

「Add & Norm」 部分的代码实现,可以直接参考 bert_layer.py#L193-L215 ,在 bert 中会循环多次使用,这里我将原始的 BertSelfOutput 和 BertOutput 和成一个了,这里的 Add & Norm 实现了三个功能:

  • Multi-Head attention 后,所有的头注意力结果是直接 concat 在一起的( view 调整 size 也可以认为 concat 在一起)直接 concat 在一起的结果用起来也有点奇怪,所以需要有个 fc ,来帮助把这些分散注意力结果合并在一起;
  • Feed Forward 操作后,纬度被提升到 intermediate_sizeBertAddNorm 还实现了把纬度从 intermediate_size 降回 hidden_size 的功能;一般来说,intermediate_sizehidden_size 的 4倍大小,非常像卷积核大小为 1 & 多个卷积核 时的情况,都是对原始输入进行放大然后再缩小,我认为可以更好的关注的输入内容的不同角度。 但 BertAddNorm 这里的实现要比卷积操作高效很多。
  • 真正的 Add & Norm 部分,也就是 layer_norm(hidden_states + input_tensor) 这一行,也就是这里的代码有多实现 densedropout 后边会有说明;
1
2
3
4
5
6
7
8
9
10
11
12
class BertAddNorm(nn.Module):
def __init__(self, intermediate_size, hidden_size, hidden_dropout_prob, layer_norm_eps):
super(BertAddNorm, self).__init__()
self.dense = nn.Linear(intermediate_size, hidden_size)
self.layer_norm = BertLayerNorm(hidden_size, eps=layer_norm_eps)
self.dropout = nn.Dropout(hidden_dropout_prob)

def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = self.dropout(hidden_states)
hidden_states = self.layer_norm(hidden_states + input_tensor)
return hidden_states

4、Feed Forward

「Position-wise Feed-Forward Networks 」 的代码实现,来自于 bert_layer.py#L218-L237

1
2
3
4
5
6
7
8
9
10
11
class BertIntermediate(nn.Module):

def __init__(self, hidden_size, intermediate_size, hidden_act):
super(BertIntermediate, self).__init__()
self.dense = nn.Linear(hidden_size, intermediate_size)
self.intermediate_act_fn = activations[hidden_act]

def forward(self, hidden_states):
hidden_states = self.dense(hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states)
return

大家可能会发现,诶? 这里怎么只有 FFN 的左半部分,外边的那个 dense 呢? 外边的那个 dense 在 Add&Norm 里边了,其实我觉得这块不太合理的,但不太好修改结构,因为修改了结构原始的参数就加载不上了。

5、Bert Layer

至此,我们可以组装出 2+3 部分,也就是 N* 循环内的下半部分,bert_layer.py#L240-L263

1
2
3
4
5
6
7
8
9
10
11
12
13
class BertAttention(nn.Module):
def __init__(self, config):
super(BertAttention, self).__init__()
self.self = MultiHeadAttentionLayer(config)
# 这里是左下的那个 Add & Norm
self.output = BertAddNorm(config.hidden_size, config.hidden_size,
config.hidden_dropout_prob, config.layer_norm_eps)

def forward(self, input_tensor, attention_mask=None, head_mask=None):
self_outputs = self.self(input_tensor, input_tensor, input_tensor, attention_mask, head_mask)
attention_output = self.output(self_outputs[0], input_tensor)
outputs = (attention_output,) + self_outputs[1:]
return outputs

并进一步得到完整的一个 bert_layerbert_layer.py#L266-L289

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class BertLayer(nn.Module):
def __init__(self, config):
super(BertLayer, self).__init__()
self.attention = BertAttention(config)

self.intermediate = BertIntermediate(config.hidden_size, config.intermediate_size, config.hidden_act)
self.output = BertAddNorm(config.intermediate_size, config.hidden_size,
config.hidden_dropout_prob, config.layer_norm_eps)

def forward(self, hidden_states, attention_mask=None, head_mask=None):
attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
attention_output = attention_outputs[0]

# 这里是左上的 Add & Norm,从而得到完整的 FFN
intermediate_output = self.intermediate(attention_output)
layer_output = self.output(intermediate_output, attention_output)

# attention_outputs[0] 是 embedding, [1] 是 attention_probs
outputs = (layer_output,) + attention_outputs[1:]
return outputs

6、Bert Encoder

将 Bert Layer 的结果,循环num_hidden_layers次,将上一轮的输出,输入到新的一轮中,代码实现 bert_model.py#L18-L52

7、Bert Pooler

对于 CLS 位,我们会进行一个特殊的 pooler 操作,即 bert_model.py#L55-L66,所以我们直接取 cls 位的结果,并不是真的第一个位置上的 embedding,而且该 embedding 经过变形并激活后的结果。

1
2
3
4
5
6
7
8
9
10
11
12
class BertPooler(nn.Module):
def __init__(self, config):
super(BertPooler, self).__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.activation = nn.Tanh()

def forward(self, hidden_states):
# 只取出第一个 token 也就是 cls 位置上的 embedding 进行 dense 变形
first_token_tensor = hidden_states[:, 0]
pooled_output = self.dense(first_token_tensor)
pooled_output = self.activation(pooled_output)
return pooled_output

8、Bert Module

这里基本上就是进行一系列合并,将 Bert Embedding 的结果输入到 BertEncoder中,具体实现 bert_model.py#L69-L185

需要注意的是,key 的替代操作,这里是因为 tf 的权重和 pytorch 权重的名称不太一样,特别是 layer_norm 的,tf 中的命名感觉不太规范,将对象命名成为了大驼峰,所以不 replace 的话就无法加载进来了。

总结

1、安装库

pip install pure_attention==0.0.20 或者 git clone link 到本地

2、下载预训练模型

这里我弄了 transformers国内下载镜像,关于 lfs,可以参考 git lfs

模型名称 git clone 自行下载
bert-base-chinese git clone git@e.coding.net:mmmwhy/file/bert-base-chinese.git https://mmmwhy.coding.net/public/file/bert-base-chinese/git/files
chinese-roberta-wwm-ext git clone git@e.coding.net:mmmwhy/file/chinese-roberta-wwm-ext.git https://mmmwhy.coding.net/public/file/chinese-roberta-wwm-ext/git/files
chinese-roberta-wwm-ext-large git lfs clone git@e.coding.net:mmmwhy/file/chinese-roberta-wwm-ext-large.git https://mmmwhy.coding.net/public/file/chinese-roberta-wwm-ext-large/git/files
ernie 1.0 git clone git@e.coding.net:mmmwhy/file/ernie-1.0.git https://mmmwhy.coding.net/public/file/ernie-1.0/git/files

速度还是比较可观的,

img

3、使用 demo

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from pure_attention.common.nlp.tokenization import Tokenizer
from pure_attention.backbone_bert.bert_model import BertModel

bert_model_path = "/data/pretrain_modal/bert-base-chinese"
test_query = "结果一致性验证"

tokenizer = Tokenizer(bert_model_path + "/vocab.txt")
bert = BertModel(bert_model_path)

tokenizer_output= tokenizer.encode(test_query, max_len=64)

our_bert_pooler_output = bert(
input_ids=tokenizer_output.input_ids,
token_type_ids=tokenizer_output.token_type_ids,
attention_mask=tokenizer_output.attention_mask).pooler_output

bert_last_hidden_state = bert(
input_ids=tokenizer_output.input_ids,
token_type_ids=tokenizer_output.token_type_ids,
attention_mask=tokenizer_output.attention_mask).last_hidden_state

4、一致性校验

在 4 种常见中文 bert 上进行实验,结果与 transformers 完全一致。校验代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import torch
from transformers import BertModel
from transformers import BertTokenizer



bert_model_path = "/data/pretrain_modal/chinese-roberta-wwm-ext-large"
test_query = "结果一致性验证"

text_tokenizer = BertTokenizer.from_pretrained(bert_model_path, do_lower_case=True)
bert_model = BertModel.from_pretrained(bert_model_path)

tensor_caption = text_tokenizer(test_query, return_tensors="pt", padding='max_length', truncation=True,
max_length=64)


origin_bert_pooler_output = bert_model(
input_ids=tensor_caption.input_ids,
attention_mask=tensor_caption.attention_mask,
token_type_ids=tensor_caption.token_type_ids).pooler_output

# 我们简化重构后的代码
from pure_attention.common.nlp.tokenization import Tokenizer as LocalTokenizer
from pure_attention.backbone_bert.bert_model import BertModel as OurBertModel
tokenizer = LocalTokenizer(bert_model_path + "/vocab.txt")
bert = OurBertModel(bert_model_path)
tokenizer_output = tokenizer.encode(test_query, max_len=64)

our_bert_pooler_output = bert(
input_ids=tokenizer_output.input_ids,
token_type_ids=tokenizer_output.token_type_ids,
attention_mask=tokenizer_output.attention_mask).pooler_output


print("check result:", torch.cosine_similarity(origin_bert_pooler_output, our_bert_pooler_output))

当时截图的时候比较早,代码稍微做了一些调整就没有新截图了,以代码为准。

5、其他部分

我一直想细致的了解一下底层代码的实现,特别是可以和 transformer 的设计图对应起来。在看了一些已有的代码后,发现 transformers 为了适应非常多种模型结构,结构变得非常复杂,代码看来比较复杂。

因此希望自己可以完成一个这样的作品,让其可以在 cv 任务和 nlp 任务上均取到 sota 的效果,我将其称之为 pure_attention

我在参考transformersbert4pytorchRead_Bert_Code的代码基础上,对结构进行了一些调整,提高了代码的易读性,并和 transformers 的结果完全一致。

pytorch 实现 bert,附带详细的注释和 transformers 国内下载链接

https://iii.run/archives/26e376bf8cc2.html

作者

mmmwhy

发布于

2022-01-23

更新于

2022-02-15

许可协议

评论