Pytorch实现Bert模型

Pytorch实现Bert模型

一、数据预处理

处理的文本是自定义的对话文本，模仿R和J两人对话

‘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

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 '.', ',', '?', '!'
# 将预处理后的词语（token）放在word_list
word_list = list(set(" ".join(sentences).split())) # ['hello', 'how', 'are', 'you',...]
# 将每一个token映射为索引 特殊编码 [PAD]：一句话的长度不够，后面添加0
word2idx = {'[PAD]' : 0, '[CLS]' : 1, '[SEP]' : 2, '[MASK]' : 3}
for i, w in enumerate(word_list): # 其他的编码通过for循环添加 从4开始，0-3已经使用了
    word2idx[w] = i + 4
# print(word2idx)
# 实现索引到token的转换
idx2word = {i: w for i, w in enumerate(word2idx)}
print(idx2word)
vocab_size = len(word2idx)
# print(vocab_size) # 没有重复词
token_list = list()
for sentence in sentences:
    arr = [word2idx[s] for s in sentence.split()]
    token_list.append(arr)
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
36
37
38
39

# 9行文本对应9个列表，每个列表中的元素个数对应于每行文本的token数
print(token_list)
1
2

二、定义Bert模型的参数

# BERT Parameters
maxlen = 30 # 每一个样本也就是每一个句子的最大长度为30，超过30要切除，少于30 补PAD
batch_size = 6
max_pred = 5 # max tokens of prediction 最多5个token被mask，虽然按照比例需要15%的token需要被mask，但是太多了设置一个上限
n_layers = 6 # encoder的层数
n_heads = 12 # multi-head的个数
d_model = 768 #word embedding 、positional embedding 、segment embedding都是768相同的维度
d_ff = 768*4 # 4*d_model, FeedForward dimension 在全连接神经网络中提升的维度 768*4 = 3072
d_k = d_v = 64  # dimension of K(=Q), V 
n_segments = 2 # 一个batch由2句话构成
1
2
3
4
5
6
7
8
9
10

三、预处理部分

拼接特殊标签，以及完成两个任务：语言模型任务以及NSP任务

# sample IsNext and NotNext to be same in small batch size
def make_data():
    batch = [] 
    positive = negative = 0 # positive若样本中两条样本相邻 加 1 ； 不相邻则negative 加 1 但最终需要保证positive 与negative比例相等
    while positive != batch_size/2 or negative != batch_size/2:
        # tokens_a_index：第一个文本的索引 tokens_b_index:第二个文本的索引 
        # randrange(0-8) 随机抽取两个索引，判断tokens_a_index + 1 是否等于 tokens_b_index 得到两个文本是否相邻
        tokens_a_index, tokens_b_index = randrange(len(sentences)), randrange(len(sentences)) # sample random index in sentences
        # 通过文本的索引tokens_a_index，获取文本中的每个token的索引放到tokens_a中；tokens_b_index同理
        tokens_a, tokens_b = token_list[tokens_a_index], token_list[tokens_b_index]
        # 在token前后拼接[CLS] 与 [SEP]
        input_ids = [word2idx['[CLS]']] + tokens_a + [word2idx['[SEP]']] + tokens_b + [word2idx['[SEP]']]
        # segment_ids 表示模型中的segment embedding，前一句全为0，个数是【 1 + len(tokens_a) + 1 】个 ，前后两个1表示特殊标识的1；后一句全0
        segment_ids = [0] * (1 + len(tokens_a) + 1) + [1] * (len(tokens_b) + 1)

        # MASK LM
        # 先取整个句子长度的15%做mask,但需要注意的是有时整个句子长度太短，比如6个token，6*0.15小于1，此时需要和1进行比较取最大值
        # 但有时句子过长，我们设置的界限是mask不超过5个，因此要和max_pred取最小值
        n_pred =  min(max_pred, max(1, int(len(input_ids) * 0.15))) # 15 % of tokens in one sentence
        # cand_maked_pos：候选mask标记，由于特殊标记[CLS]和[SEP]做mask是无意义的，因此需要排除
        cand_maked_pos = [i for i, token in enumerate(input_ids)
                          # 只要不是[CLS]和[SEP] 就可以将索引存入cand_maked_pos
                          if token != word2idx['[CLS]'] and token != word2idx['[SEP]']] # candidate masked position
        shuffle(cand_maked_pos) # 由于是随机mask,将cand_maked_pos列表中的元素随机打乱
        masked_tokens, masked_pos = [], []
        for pos in cand_maked_pos[:n_pred]: # 取乱序的索引cand_maked_pos前n_pred 做mask
            masked_pos.append(pos) # masked_pos：mask标记对应的索引
            masked_tokens.append(input_ids[pos]) # masked_tokens：mask标记对应的原来的token值
            # Bert模型中mask标记有80%被替换为真正的mask，10% 被随机替换为词表中的任意词，10%不变
            if random() < 0.8:  # 80% 的 概率 被替换为真正的mask
                input_ids[pos] = word2idx['[MASK]'] # make 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' 但不能是特殊标记，也就是说索引要大于4
                  index = randint(0, vocab_size - 1) # 索引小于4需要重新获取一个随机数
                input_ids[pos] = index # replace 用随机的词替换该位置的token

        # Zero Paddings
        # 对长度不足maxlen30的文本 补 PAD
        n_pad = maxlen - len(input_ids) # 30 - 文本长度 = 需要补 0 的个数
        input_ids.extend([0] * n_pad) # 不足30的位置token补0
        segment_ids.extend([0] * n_pad) # segment embedding 补 0 

        # Zero Padding (100% - 15%) tokens
        if max_pred > n_pred: 
            n_pad = max_pred - n_pred
            masked_tokens.extend([0] * n_pad) # 保证masked_tokens和masked_pos长度始终为max_pred(5)
            masked_pos.extend([0] * n_pad)

        # 判断两个文本是否相邻  tokens_a_index + 1 是否等于 tokens_b_index  ；前提是positive 和negative比例相等
        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
        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
    print(batch)
    return batch
# Proprecessing Finished 数据预处理结束

batch = make_data()
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)
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

四、Bert模型构建

def get_attn_pad_mask(seq_q, seq_k):
    batch_size, seq_len = seq_q.size()
    # eq(zero) is PAD token
    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]

# bert论文中提出的新激活函数gelu()
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__()
        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
        self.norm = nn.LayerNorm(d_model)

    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) # 最终的embedding为三者相加 三者维度都是相等的
        return self.norm(embedding)

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()
    # 通过 Q 和 K 计算出 scores，然后将 scores 和 V 相乘，得到每个单词的 context vector
    # Q: [batch_size, seq_len, d_model], K: [batch_size, seq_len, d_model], V: [batch_size, seq_len, d_model]
    def forward(self, Q, K, V, attn_mask):
        # transpose(-1, -2) 相当于将K矩阵转置
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # scores : [batch_size, n_heads, seq_len, seq_len]
        # 相乘之后得到的 scores 还不能立刻进行 softmax，需要和 attn_mask 相加，把一些需要屏蔽的信息屏蔽掉，填充为-1e9 也就是负无穷
        # attn_mask 是一个仅由 True 和 False 组成的 tensor，并且一定会保证 attn_mask 和 scores 的维度四个值相同（不然无法做对应位置相加）
        scores.masked_fill_(attn_mask, -1e9) # Fills elements of self tensor with value where mask is one.
        attn = nn.Softmax(dim=-1)(scores)
        context = torch.matmul(attn, V) # Softmax后的值与 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):
        # Q：第一个enc_inputs，K：第二个enc_inputs，V：第三个enc_inputs
        # 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(H,S进行转置)-> (B, H, S, W)
        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]

        # 将attn_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]
        context = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
        # context.transpose(1, 2) :[batch_size, seq_len, n_heads, d_v]
        # contiguous() : 把tensor变成在内存中连续分布的形式  四维->三维
        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)# 将维度变为d_model
        return nn.LayerNorm(d_model)(output + residual) # output: [batch_size, seq_len, d_model] 残差连接 后做正交化

class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        # 维度转换 d_model<-> d_ff
        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() # 前馈神经网络

    def forward(self, enc_inputs, enc_self_attn_mask):
        # self-attention 通过MultiHeadAttention实现 传入三个enc_inputs作用是分别于W(Q,K,V)相乘生成 Q,K,V矩阵
        enc_outputs = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask) # enc_inputs to same Q,K,V
        # pos_ffn 特征提取 通过PoswiseFeedForwardNet实现
        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() # 构建词向量矩阵
        # 使用ModuleList对6个encoder进行堆叠
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])  # n_layers = 6 # encoder的层数 
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_model), # 添加全连接层，维度没有发生变化 cls二分类任务
            nn.Dropout(0.5),
            nn.Tanh(), # 激活函数 Tanh()
        )
        self.classifier = nn.Linear(d_model, 2) # cls最后的输出，从d_model-> 2维，也就是二分类任务
        self.linear = nn.Linear(d_model, d_model) # 添加全连接层，维度没有发生变化 mlm任务
        self.activ2 = gelu # 激活函数2 gelu()
        
        # fc2 is shared with embedding layer
        embed_weight = self.embedding.tok_embed.weight
        self.fc2 = nn.Linear(d_model, vocab_size, bias=False) # d_model-> vocab_size 解码
        self.fc2.weight = embed_weight

    # encoder的输入包括三个部分 
    def forward(self, input_ids, segment_ids, masked_pos): 
        # 生成input_ids对应的embedding 以及 segment_ids对应的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]
            output = layer(output, enc_self_attn_mask)
        # it will be decided by first token(CLS) 取第一个cls ，fc里面包括了前馈神经网络层 以及激活函数 
        h_pooled = self.fc(output[:, 0]) # [batch_size, d_model]
        logits_clsf = self.classifier(h_pooled) # [batch_size, 2] predict isNext 在最后过二分类得到预测输出

        # 例如masked_pos = [6,5,17,0,0] 也就是说第6，5，17个被mask
        masked_pos = masked_pos[:, :, None].expand(-1, -1, d_model) # [batch_size, max_pred, d_model]
        # 根据masked_pos利用gather()在output中取出第6，5，17个token 
        # 将token 与masked_pos进行对齐（因为一开始token是按照[0,1,2...]排列的） 要将token转变成[6,5,17,...]才能对齐
        h_masked = torch.gather(output, 1, masked_pos) # masking position [batch_size, max_pred, d_model]
        # 让h_masked(带mask的token)过全连接层 ，再过激活函数2 
        h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, d_model]
        # 再解码到词表大小 d_model-> vocab_size
        logits_lm = self.fc2(h_masked) # [batch_size, max_pred, vocab_size]
        return logits_lm, logits_clsf # logits_lm计算mask对应的词 logits_clsf计算前后两个句子是否连续
model = BERT()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adadelta(model.parameters(), lr=0.001)
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

五、torch.gather()函数实现的功能

out = torch.gather(input, dim, index)
# out[i][j][k] = input[index[i][j][k]][j][k] # dim = 0
# out[i][j][k] = input[i][index[i][j][k]][k] # dim = 1
# out[i][j][k] = input[i][j][index[i][j][k]] # dim = 2
1
2
3
4

输入的数据以三维为例：
当dim = 0 时，输出的第一维被替换为index[i][j][k]，第二维和第三维不变；
当dim = 1 时，输出的第二维被替换为index[i][j][k]，第一维和第三维不变；
当dim = 2 时，输出的第三维被替换为index[i][j][k]，第一维和第二维不变。

以具体向量为例：

index = torch.from_numpy(np.array([[1, 2, 0], [2, 0, 1]])).type(torch.LongTensor)
index = index[:, :, None].expand(-1, -1, 10)
print(index)
1
2
3

首先定义两个tensor：[1，2，0]、[2，0，1]，将其拓展为三维，每一维长度为10

接着随机生成一个 [2, 3, 10] 维的 tensor，可以理解为有 2 个 batch，每个 batch 有 3 句话，每句话由 10 个词构成，只不过这里的词不是以正整数（索引）的形式出现，而是连续的数值。

input = torch.rand(2, 3, 10)
print(input)
1
2

调用 torch.gather(input, 1, index) 函数

print(torch.gather(input, 1, index))
1

index 中第一行的 tensor 会作用于 input 的第一个 batch，具体来说，原本三句话的顺序是 [0, 1, 2]，现在会根据 [1, 2, 0] 调换顺序。index 中第 2 行的 tensor 会作用于 input 的第二个 batch，具体来说，原本三句话的顺序是 [0, 1, 2]，现在会根据 [2, 0, 1] 调换顺序。

六、transpose()函数作用

实现两个矩阵转置

index = torch.from_numpy(np.array([[1, 2, 0], [2, 0, 1]])).type(torch.LongTensor)
print(index)
print(index.transpose(-1, -2))
1
2
3

七、训练 & 计算损失

for epoch in range(180):
    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
      loss_lm = (loss_lm.float()).mean()
      loss_clsf = criterion(logits_clsf, isNext) # for sentence classification
      loss = loss_lm + loss_clsf # 损失为两者之和
      if (epoch + 1) % 10 == 0:
          print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss))
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12

八、测试

# Predict mask tokens ans isNext
input_ids, segment_ids, masked_tokens, masked_pos, isNext = batch[1]
print(text)
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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14