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)

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

二、定义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句话构成

三、预处理部分

拼接特殊标签，以及完成两个任务：语言模型任务以及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；后一句全0segment_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_posif token != word2idx['[CLS]'] and token != word2idx['[SEP]']] # candidate masked positionshuffle(cand_maked_pos) # 由于是随机mask,将cand_maked_pos列表中的元素随机打乱masked_tokens, masked_pos = [], []for pos in cand_maked_pos[:n_pred]: # 取乱序的索引cand_maked_pos前n_pred 做maskmasked_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% 的 概率 被替换为真正的maskinput_ids[pos] = word2idx['[MASK]'] # make maskelif random() > 0.9:  # 10% 的概率被随机替换为词表中的任意词index = randint(0, vocab_size - 1) # random index in vocabulary 从词表中随机选择一个词的索引 可以是本身while index < 4: # can't involve 'CLS', 'SEP', 'PAD' 但不能是特殊标记，也就是说索引要大于4index = randint(0, vocab_size - 1) # 索引小于4需要重新获取一个随机数input_ids[pos] = index # replace 用随机的词替换该位置的token# Zero Paddings# 对长度不足maxlen30的文本 补 PADn_pad = maxlen - len(input_ids) # 30 - 文本长度 = 需要补 0 的个数input_ids.extend([0] * n_pad) # 不足30的位置token补0segment_ids.extend([0] * n_pad) # segment embedding 补 0# Zero Padding (100% - 15%) tokensif max_pred > n_pred: n_pad = max_pred - n_predmasked_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]) # IsNextpositive += 1elif tokens_a_index + 1 != tokens_b_index and negative < batch_size/2:batch.append([input_ids, segment_ids, masked_tokens, masked_pos, False]) # NotNextnegative += 1print(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_idsself.segment_ids = segment_idsself.masked_tokens = masked_tokensself.masked_pos = masked_posself.isNext = isNextdef __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)

四、Bert模型构建

def get_attn_pad_mask(seq_q, seq_k):batch_size, seq_len = seq_q.size()# eq(zero) is PAD tokenpad_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 embeddingself.pos_embed = nn.Embedding(maxlen, d_model)  # position embeddingself.seg_embed = nn.Embedding(n_segments, d_model)  # segment(token type) embeddingself.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 contextclass 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_modelreturn 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_ffself.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_outputsclass 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 layerembed_weight = self.embedding.tok_embed.weightself.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对应的embeddingoutput = 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个被maskmasked_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)过全连接层 ，再过激活函数2h_masked = self.activ2(self.linear(h_masked)) # [batch_size, max_pred, d_model]# 再解码到词表大小 d_model-> vocab_sizelogits_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)

五、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

输入的数据以三维为例：
当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)

首先定义两个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)

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

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

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))

七、训练 & 计算损失

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 LMloss_lm = (loss_lm.float()).mean()loss_clsf = criterion(logits_clsf, isNext) # for sentence classificationloss = 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()

八、测试

# 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)