在pytorch上实现bert的简单预训练过程

前言:博主是一名研一在读学生,刚刚接触nlp不久,作品如有纰漏之处,欢迎大家批评指正,谢谢!(另外本文代码不是自己原创,解释和思路为原创。文章创作目的在于分享和知识复习,无任何盈利目的)

本文包括原理和代码设计思路部分,数据预处理部分,模型部分和训练验证部分四大块,建议阅读时间20分钟。(后附完整代码)

一、代码设计思想

大家都知道,bert是谷歌在超大数据集上通过无监督学习,学习包含字间关系的字或词的特征表示。本文篇幅有限,不可能将bert在超大参数上的完整训练过程讲清楚。博主是个菜鸡,完整的过程不会。所以本文会在一个简单的数据集上做bert的预训练。但是麻雀虽小,五脏俱全。

先来介绍预处理思路。bert训练主要包括两个部分,分别是token级别的MLM和sentence级别的next sentence。所以我们对句子的处理包括将两句话放在一起去预测是否是上下句关系,以及将句子中词随机mask掉,并让计算机去预测这些词是什么。而计算机只能识别数字,不能识别人类语言,所以呢,为了满足这些条件,我们的模型输入就有了。就是一个由两句话组成,这两句话中词部分会被mask掉,而且每个词都是由数字来表示的input_ids。两个任务当然就有两个loss。对于sentence,我们需要知道输入是否是上下句关系的正确结果isNext,对于MLM,我们要知道被MASK的词在input_ids的位置masked_pos以便从模型的输出结果中提取出对应位置的预测的MASK的词去与正确的MASK的词masked_tokens做loss。另外为了区别哪些词是第一句话哪些是第二句话,我们还需要加入segment_ids(将第一句话标为1,将第二句话标为2)。

总结,将原始文本预处理后我们需要得到五个值:inputs_ids,segment_ids,masked_tokens,masked_pos,is_Next

再来聊聊bert模型的设计思路。原理就不多聊了,网上多的汗牛充栋。思路就是将输入做一个embedding后,乘上初始化的QKV三个矩阵,得到相应的q_s,k_s,v_s。q与k做内积得到分数,再去与v做内积后得到与输入维度相同的输出,再与输入做残差并归一化,得到最终的输出。如果attention是多层的,就把最后的输出重新放入模型的输入继续训练。没听明白没关系,这部分会在代码部分详细解释,现在有个大致思路就行:input—>embedding—>QKV–(加上embedding后的input)->output。输出通过各种变换与is_Next和masked_tokens做loss。然后反向传播,更新权重,降低loss。

二、数据预处理

先来看原始数据:

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

为了构建inputs_ids就要构建每句话的idx表示(token_list),为了构建token_list就要构建字典(word2idx)

第一二句的意思是将句子中的符号去掉,转化为小写,用’\n’做分割转化为list。接着将list以空格为间隔在转化为str,再将该str以空格为间隔返回列表,做set去重,得到词列表。

在字典中加入了4个基本的辅助token,分别是表示填充的[PAD],表示语句开头的[CLS],表示句子间隔的[SEP],表示遮蔽词汇的[MASK]。

sentences=re.sub('[.,?]','',text.lower()).split('\n')
word_list=list(set(' '.join(sentences).split()))
word2idx={'[PAD]':0,'[CLS]':1,'[SEP]':2,'[MASK]':3}
for i,w in enumerate(word_list):
    word2idx[w]=i+4
idx2word={i:w for i,w in enumerate(word2idx)}
vocab_size=len(word2idx)

通过word2idx构建token_list:

token_list=[]
for sen in sentences:
    token_list.append([word2idx[i] for i in sen.split()])

先来一段参数设置:

maxlen=30 #inputs_idsd的最大长度(最长两句话的长度为25。这里设成30)
batch_size=6
max_pred=5 #max mask
n_layers=1
n_heads=12 #模型中attention一共有12层
d_model=768
d_ff=768*4 #模型中全连接神经网络的维度
d_k=d_v=64
n_segments=2 #每一行由多少句话构成

接下来是大头,是重点也是难点,如何通过token_list得到bert所需要的inputs_ids,segment_ids,masked_tokens,masked_pos,is_Next。

在bert的训练过程中,为了保证结果的有效性,需要保证样本类别数量一致。所以在beachsize为6的样本下我们需要让isNext和not isNext的个数都为3个。mask的词的数量是句子长度的15%。但是根据本例实际我们不能让他的数量小于1,也不能大于5。mask的词中80%要被替换为[mask],10%要被替换为词库中其他的随机词,10%不变。所以我们接下来一起跟着代码来看如何利用random包实现上述所讲。

(解释:cand_maked_pos是input_ids中所有可以被mask的词的位置,将其shuffle打乱后取前n_pred个,就取到了被mask的词的位置)

#IsNext和NotNext的个数得一样
def make_data():
    batch=[]
    positive=negative=0
    while positive != batch_size/2 or negative != batch_size/2:
        tokens_a_index, tokens_b_index = randrange(len(sentences)),randrange(len(sentences))
        tokens_a,tokens_b=token_list[tokens_a_index],token_list[tokens_b_index]
        input_ids=[word2idx['[CLS]']]+tokens_a+[word2idx['[SEP]']]+tokens_b+[word2idx['[SEP]']]
        segment_ids=[0]*(1+len(tokens_a)+1)+[1]*(len(tokens_b)+1)

        #MASK LM
        n_pred=min(max_pred,max(1,int(len(input_ids)*0.15))) #15% of tokens in one sentence
        #被mask的值不能是cls和sep:
        cand_maked_pos=[i for i,token in enumerate(input_ids)
                        if token != word2idx['[CLS]'] and token != word2idx['[SEP]']]
        shuffle(cand_maked_pos)
        masked_tokens,masked_pos=[],[]
        for pos in cand_maked_pos[:n_pred]:
            #将被mask位置的正确的位置和值保存下来
            masked_pos.append(pos)
            masked_tokens.append(input_ids[pos])
            #在原数据(input_ids)上替换为mask或者其他词
            if random()0.9:
                index=randint(4,vocab_size-1)
                #用词库中的词来替换,但是不能用cls,sep,pad,mask来替换
                # while indexn_pred:
            n_pad=max_pred-n_pred
            masked_tokens.extend([0]*n_pad)
            masked_pos.extend([0]*n_pad)

        #需要确保正确样本数和错误样本数一样
        if tokens_a_index+1==tokens_b_index and positive < batch_size/2:
            batch.append([input_ids,segment_ids,masked_tokens,masked_pos,True])
            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])
            negative+=1
"""
    此时batch长度为6,每个长度中有四个量,分别是:
    input_ids:含有两句话以及cls,sep的id(其中的值已经被mask替换过了)。因为做了pad,此时batch中不同input_ids的长度相同
    segment_ids:0表示第一句话,1表示第二句话
    masked_tokens:保存的被替换的词,用于做loss
    masked_pos:保存的被替换的词在input_ids中的位置。
"""
    return batch

构建dataset和dataloader:

batch=make_data()
input_ids,segment_ids,masked_token,masked_pos,isNext=zip(*batch)
input_ids,segment_ids,masked_token,masked_pos,isNext=torch.LongTensor(input_ids),\
    torch.LongTensor(segment_ids),torch.LongTensor(masked_token),torch.LongTensor(masked_pos),\
    torch.LongTensor(isNext)

class MyDataSet(Dataset):
    def __init__(self,input_ids,segment_ids,masked_token,masked_pos,isNext):
        self.input_ids=input_ids
        self.segment_ids=segment_ids
        self.masked_token=masked_token
        self.masked_pos=masked_pos
        self.isNext=isNext

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, item):
        return self.input_ids[item],self.segment_ids[item],self.masked_token[item],self.masked_pos[item],self.isNext[item]

trainset=MyDataSet(input_ids,segment_ids,masked_token,masked_pos,isNext)
trainloader=DataLoader(dataset=trainset,batch_size=batch_size,shuffle=True)

至此,数据预处理就介绍完了。(是不是感觉还挺复杂的,其实数据本身很简单,只是模型是双任务,所以处理起来就相对麻烦一些)

三、bert模型

1、这部分会更复杂,我们先跟着bert中的forword看一下数据在模型中的流动过程,都经过了哪些处理

    def forward(self,input_ids,segment_ids,masked_pos):
        output=self.embedding(input_ids,segment_ids) #[batch_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=layer(output,enc_self_attn_mask)
        #it will be decided by first token(CLS)
        h_pooled=self.fc(output[:,0]) #[batch_size,d_model]。且对于三维tensor,[:,0]和[:,0,:]效果是一样的
        #且[]中出现了数字,他就是降维操作,所以output由三维降到两维。
        logits_clsf=self.classifier(h_pooled)
        #得到判断两句话是不是上下句关系的结果

        #得到被mask位置的词,准备与正确词进行比较
        masked_pos=masked_pos[:,:,None].expand(-1,-1,d_model) #[batch_size,max_pred,d_model]
        h_masked=torch.gather(output,1,masked_pos) #masking position [batch_size,max_pred,d_model]
        h_masked=self.activ2(self.linear(h_masked))# [batch_size,max_pred,d_model]
        logits_lm=self.fc2(h_masked) #[batch_size,max_pred,vocab_size]
        return logits_lm,logits_clsf #预测的被mask地方的值,预测两句话是否为上下句的结果

首先,输入需要经过embedding从二维变成三维。之前我们对很多句子长度不足maxlen的做了补零操作,而补零这部分是不需要参与运算的,enc_self_attn_mask的作用就是屏蔽补0部分。接着layer就是attention的处理层,这里放了一个循环,数据以输入输入都是(barchsize,maxlen,d_model)的维度在layer层中循环了6次,得到了最终表征每句话的词向量。这时候取每个batch的第0维,就是cls的向量为h_pooled准备做isNext判断。masked_pos保存的是被mask词的位置,用gather函数取得output中masked_pos位置的词,准备做mask词预测。

2、好了,接着来看每一步是怎么做的:

下图是embedding层,将输入词信息,位置信息,是否是第一句话的信息全部做嵌入后做信息融合。得到最终的三维输出(batchsize,maxlen,d_model)

class Embedding(torch.nn.Module):
    def __init__(self):
        super(Embedding, self).__init__()
        self.tok_embed=torch.nn.Embedding(vocab_size,d_model)
        self.pos_embed=torch.nn.Embedding(maxlen,d_model)
        self.seg_embed=torch.nn.Embedding(n_segments,d_model)
        self.norm=torch.nn.LayerNorm(d_model)
    def forward(self, x,seg):
        seq_len=x.size(1)
        pos=torch.arange(seq_len,dtype=torch.long).to(device)
        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)
        return self.norm(embedding)

3、下图就是上边提到的enc_self_attn_mask的构建过程,输入中与0相等(即做[PAD])的部分置为True,其他部分为False。

另外声明bert自带的激活函数gelu。

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) #[batchsize,1,seq_len]
    #eq(0)表示和0相等的返回True,不相等返回False。unsqueeze(1)在第一维上
    return pad_attn_mask.expand(batch_size,seq_len,seq_len) #[batchsize,seq_len,seq_len]

def gelu(x):
    return x*0.5*(1.0+torch.erf(x/math.sqrt(2.0)))

4、下图就是多头注意力机制的处理过程了,经过维度变换后,qk做乘积得到分数,分数与attn_mask(即上文提到的enc_self_attn_mask)做masked_fill_,将scores中对应attn_mask中为True的部分的值替换为-1e9,这样进入softmax后做e**x时就把这部分的分数屏蔽掉了。处理完成后经过维度变换,与输入做残差,得到最终输出。

class ScaleDotProductAttention(torch.nn.Module):
    def __init__(self):
        super(ScaleDotProductAttention, self).__init__()

    def forward(self,Q,K,V,attn_mask):
        scores=torch.matmul(Q,K.transpose(-1,-2))/np.sqrt(d_k) #matmul做内积,transpose做转置
        scores.masked_fill_(attn_mask,-1e9) #fills elements of self tensor with value where mask is one.

        attn=torch.nn.Softmax(dim=-1)(scores)
        context=torch.matmul(attn,V)
        return context
class MultiHeadAttention(torch.nn.Module):
    def __init__(self):
        super(MultiHeadAttention,self).__init__()
        self.W_Q=torch.nn.Linear(d_model,d_k*n_heads)
        self.W_K = torch.nn.Linear(d_model, d_k * n_heads)
        self.W_V = torch.nn.Linear(d_model, d_k * n_heads)
        self.linear=torch.nn.Linear(n_heads*d_v,d_model)
        self.LN=torch.nn.LayerNorm(d_model)
    def forward(self,Q,K,V,attn_mask):
        residual,batch_size=Q,Q.size(0)
        #(B,S,D) --proj-> (B,S,D) --split-> (B,S,H,W) --trans-> (B,H,S,W)
        q_s=self.W_Q(Q).view(batch_size,-1,n_heads,d_k).transpose(1,2)  #[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)  # [batch_size,n_heads,seq_len,d_k]
        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # [batch_size,n_heads,seq_len,d_k]

        attn_mask=attn_mask.unsqueeze(1).repeat(1,n_heads,1,1) #将扩展的第二维扩张到n_heads。[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=ScaleDotProductAttention().to(device)(q_s,k_s,v_s,attn_mask)
        context=context.transpose(1,2).contiguous().view(batch_size,-1,n_heads*d_v)
        output=self.linear(context)
        return self.LN(output+residual)  #这边应该是用到了残差吧
        #output: [batch_size,seq_len,d_model]

5、下图是对数据做维度变化。这部分与整体数据处理思路没啥大关系,应该是加上后可以让结果更好吧。

class PoswiseFeedForwardNet(torch.nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet,self).__init__()
        self.fc1=torch.nn.Linear(d_model,d_ff)
        self.fc2=torch.nn.Linear(d_ff,d_model)

    def forward(self, x):
        return self.fc2(gelu(self.fc1(x)))

6、下图将attention和维度变化整合在一起,这就是数据embedding后的整个编码部分。准备被bert调用。

class EncoderLayer(torch.nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn=MultiHeadAttention().to(device)
        self.pos_ffn=PoswiseFeedForwardNet().to(device)

    def forward(self, enc_inputs,enc_self_attn_mask):
        enc_outputs=self.enc_self_attn(enc_inputs,enc_inputs,enc_inputs,enc_self_attn_mask) #same Q,K,V
        enc_outputs=self.pos_ffn(enc_outputs)
        return enc_outputs

7、bert模型,loss,优化器:

class BERT(torch.nn.Module):
    def __init__(self):
        super(BERT,self).__init__()
        self.embedding=Embedding().to(device) #这是上边定义的Embedding,self.embedding就是一个类对象
        self.layers=torch.nn.ModuleList([EncoderLayer().to(device) for _ in range(n_layers)])
        self.fc=torch.nn.Sequential(
            torch.nn.Linear(d_model,d_model),
            torch.nn.Dropout(0.5),
            torch.nn.Tanh(),
        )
        self.classifier=torch.nn.Linear(d_model,2)
        self.linear=torch.nn.Linear(d_model,d_model)
        self.activ2=gelu #这是上边定义的gelu,self.activ2就是gelu函数
        #这里为什么要用embed_weight的维度,回头再看一下
        embed_weight=self.embedding.tok_embed.weight
        self.fc2=torch.nn.Linear(d_model,vocab_size,bias=False)
        self.fc2.weight=embed_weight

    def forward(self,input_ids,segment_ids,masked_pos):
        output=self.embedding(input_ids,segment_ids) #[batch_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=layer(output,enc_self_attn_mask)
        #it will be decided by first token(CLS)
        h_pooled=self.fc(output[:,0]) #[batch_size,d_model]。且对于三维tensor,[:,0]和[:,0,:]效果是一样的
        #且[]中出现了数字,他就是降维操作,所以output由三维降到两维。
        logits_clsf=self.classifier(h_pooled)
        #得到判断两句话是不是上下句关系的结果

        #得到被mask位置的词,准备与正确词进行比较
        masked_pos=masked_pos[:,:,None].expand(-1,-1,d_model) #[batch_size,max_pred,d_model]
        h_masked=torch.gather(output,1,masked_pos) #masking position [batch_size,max_pred,d_model]
        h_masked=self.activ2(self.linear(h_masked))# [batch_size,max_pred,d_model]
        logits_lm=self.fc2(h_masked) #[batch_size,max_pred,vocab_size]
        return logits_lm,logits_clsf #预测的被mask地方的值,预测两句话是否为上下句的结果

model=BERT()
criterion=torch.nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(model.parameters(),lr=0.001)

至此,模型整个部分就讲解结束了,下面还包括训练部分和测试部分,比较简单,就放在整体代码里了,大家如果跑不通或者训练的过程中遇到问题,可以给我留言。

整体代码:

import re
import torch
import math
import numpy as np
from random import *
from torch.utils.data import Dataset,DataLoader
device=torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

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 Romeol\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
)
"""
下面这段代码很有智慧,做了数据数据预处理,构建词表,构建双向字典(由idx向word,由word向idx)
在短短7行代码中完成

预处理思路:
肯定要先构建词典,(词对索引,索引对词,词典长度)
再根据词典,对应句子中(不管是一维还是二维)的词找其所对应的索引
"""
sentences=re.sub('[.,?!]','',text.lower()).split('\n')
word_list=list(set(' '.join(sentences).split()))
word2idx={'[PAD]':0,'[CLS]':1,'[SEP]':2,'[MASK]':3}
for i,w in enumerate(word_list):
    word2idx[w]=i+4
idx2word={i:w for i,w in enumerate(word2idx)}
vocab_size=len(word2idx)

#做token_list
token_list=[]
for sen in sentences:
    token_list.append([word2idx[i] for i in sen.split()]) #[[29, 6, 32, 37, 33, 31, 13], [29, 13, 16, 22, 21...

#bert parameters
maxlen=30
batch_size=6
max_pred=5 #max mask
n_layers=1
n_heads=12
d_model=768
d_ff=768*4 #全连接神经网络的维度
d_k=d_v=64
n_segments=2 #每一行由多少句话构成

#IsNext和NotNext的个数得一样
def make_data():
    batch=[]
    positive=negative=0
    while positive != batch_size/2 or negative != batch_size/2:
        tokens_a_index, tokens_b_index = randrange(len(sentences)),randrange(len(sentences))
        tokens_a,tokens_b=token_list[tokens_a_index],token_list[tokens_b_index]
        input_ids=[word2idx['[CLS]']]+tokens_a+[word2idx['[SEP]']]+tokens_b+[word2idx['[SEP]']]
        segment_ids=[0]*(1+len(tokens_a)+1)+[1]*(len(tokens_b)+1)

        #MASK LM
        n_pred=min(max_pred,max(1,int(len(input_ids)*0.15))) #15% of tokens in one sentence
        #被mask的值不能是cls和sep:
        cand_maked_pos=[i for i,token in enumerate(input_ids)
                        if token != word2idx['[CLS]'] and token != word2idx['[SEP]']]
        shuffle(cand_maked_pos)
        masked_tokens,masked_pos=[],[]
        for pos in cand_maked_pos[:n_pred]:
            #将被mask位置的正确的位置和值保存下来
            masked_pos.append(pos)
            masked_tokens.append(input_ids[pos])
            #在原数据(input_ids)上替换为mask或者其他词
            if random()0.9:
                index=randint(4,vocab_size-1)
                #用词库中的词来替换,但是不能用cls,sep,pad,mask来替换
                # while indexn_pred:
            n_pad=max_pred-n_pred
            masked_tokens.extend([0]*n_pad)
            masked_pos.extend([0]*n_pad)

        #需要确保正确样本数和错误样本数一样
        if tokens_a_index+1==tokens_b_index and positive < batch_size/2:
            batch.append([input_ids,segment_ids,masked_tokens,masked_pos,True])
            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])
            negative+=1
"""
    此时batch长度为6,每个长度中有四个量,分别是:
    input_ids:含有两句话以及cls,sep的id(其中的值已经被mask替换过了)。因为做了pad,此时batch中不同input_ids的长度相同
    segment_ids:0表示第一句话,1表示第二句话
    masked_tokens:保存的被替换的词,用于做loss
    masked_pos:保存的被替换的词在input_ids中的位置。
"""
    return batch
batch=make_data()
input_ids,segment_ids,masked_token,masked_pos,isNext=zip(*batch)
input_ids,segment_ids,masked_token,masked_pos,isNext=torch.LongTensor(input_ids),\
    torch.LongTensor(segment_ids),torch.LongTensor(masked_token),torch.LongTensor(masked_pos),\
    torch.LongTensor(isNext)

class MyDataSet(Dataset):
    def __init__(self,input_ids,segment_ids,masked_token,masked_pos,isNext):
        self.input_ids=input_ids
        self.segment_ids=segment_ids
        self.masked_token=masked_token
        self.masked_pos=masked_pos
        self.isNext=isNext

    def __len__(self):
        return len(self.input_ids)

    def __getitem__(self, item):
        return self.input_ids[item],self.segment_ids[item],self.masked_token[item],self.masked_pos[item],self.isNext[item]

trainset=MyDataSet(input_ids,segment_ids,masked_token,masked_pos,isNext)
trainloader=DataLoader(dataset=trainset,batch_size=batch_size,shuffle=True)

#这里seq_k不用吗??
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) #[batchsize,1,seq_len]
    #eq(0)表示和0相等的返回True,不相等返回False。unsqueeze(1)在第一维上
    return pad_attn_mask.expand(batch_size,seq_len,seq_len) #[batchsize,seq_len,seq_len]

def gelu(x):
    return x*0.5*(1.0+torch.erf(x/math.sqrt(2.0)))

class Embedding(torch.nn.Module):
    def __init__(self):
        super(Embedding, self).__init__()
        self.tok_embed=torch.nn.Embedding(vocab_size,d_model)
        self.pos_embed=torch.nn.Embedding(maxlen,d_model)
        self.seg_embed=torch.nn.Embedding(n_segments,d_model)
        self.norm=torch.nn.LayerNorm(d_model)
    def forward(self, x,seg):
        seq_len=x.size(1)
        pos=torch.arange(seq_len,dtype=torch.long).to(device)
        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)
        return self.norm(embedding)

#没有特别懂,需要结合下边的代码
class ScaleDotProductAttention(torch.nn.Module):
    def __init__(self):
        super(ScaleDotProductAttention, self).__init__()

    def forward(self,Q,K,V,attn_mask):
        scores=torch.matmul(Q,K.transpose(-1,-2))/np.sqrt(d_k) #matmul做内积,transpose做转置
        scores.masked_fill_(attn_mask,-1e9) #fills elements of self tensor with value where mask is one.

        attn=torch.nn.Softmax(dim=-1)(scores)
        context=torch.matmul(attn,V)
        return context

class MultiHeadAttention(torch.nn.Module):
    def __init__(self):
        super(MultiHeadAttention,self).__init__()
        self.W_Q=torch.nn.Linear(d_model,d_k*n_heads)
        self.W_K = torch.nn.Linear(d_model, d_k * n_heads)
        self.W_V = torch.nn.Linear(d_model, d_k * n_heads)
        self.linear=torch.nn.Linear(n_heads*d_v,d_model)
        self.LN=torch.nn.LayerNorm(d_model)
    def forward(self,Q,K,V,attn_mask):
        residual,batch_size=Q,Q.size(0)
        #(B,S,D) --proj-> (B,S,D) --split-> (B,S,H,W) --trans-> (B,H,S,W)
        q_s=self.W_Q(Q).view(batch_size,-1,n_heads,d_k).transpose(1,2)  #[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)  # [batch_size,n_heads,seq_len,d_k]
        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # [batch_size,n_heads,seq_len,d_k]

        attn_mask=attn_mask.unsqueeze(1).repeat(1,n_heads,1,1) #将扩展的第二维扩张到n_heads。[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=ScaleDotProductAttention().to(device)(q_s,k_s,v_s,attn_mask)
        context=context.transpose(1,2).contiguous().view(batch_size,-1,n_heads*d_v)
        output=self.linear(context)
        return self.LN(output+residual)  #这边应该是用到了残差吧
        #output: [batch_size,seq_len,d_model]

class PoswiseFeedForwardNet(torch.nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet,self).__init__()
        self.fc1=torch.nn.Linear(d_model,d_ff)
        self.fc2=torch.nn.Linear(d_ff,d_model)

    def forward(self, x):
        return self.fc2(gelu(self.fc1(x)))

class EncoderLayer(torch.nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn=MultiHeadAttention().to(device)
        self.pos_ffn=PoswiseFeedForwardNet().to(device)

    def forward(self, enc_inputs,enc_self_attn_mask):
        enc_outputs=self.enc_self_attn(enc_inputs,enc_inputs,enc_inputs,enc_self_attn_mask) #same Q,K,V
        enc_outputs=self.pos_ffn(enc_outputs)
        return enc_outputs

class BERT(torch.nn.Module):
    def __init__(self):
        super(BERT,self).__init__()
        self.embedding=Embedding().to(device) #这是上边定义的Embedding,self.embedding就是一个类对象
        self.layers=torch.nn.ModuleList([EncoderLayer().to(device) for _ in range(n_layers)])
        self.fc=torch.nn.Sequential(
            torch.nn.Linear(d_model,d_model),
            torch.nn.Dropout(0.5),
            torch.nn.Tanh(),
        )
        self.classifier=torch.nn.Linear(d_model,2)
        self.linear=torch.nn.Linear(d_model,d_model)
        self.activ2=gelu #这是上边定义的gelu,self.activ2就是gelu函数
        #这里为什么要用embed_weight的维度,回头再看一下
        embed_weight=self.embedding.tok_embed.weight
        self.fc2=torch.nn.Linear(d_model,vocab_size,bias=False)
        self.fc2.weight=embed_weight

    def forward(self,input_ids,segment_ids,masked_pos):
        output=self.embedding(input_ids,segment_ids) #[batch_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=layer(output,enc_self_attn_mask)
        #it will be decided by first token(CLS)
        h_pooled=self.fc(output[:,0]) #[batch_size,d_model]。且对于三维tensor,[:,0]和[:,0,:]效果是一样的
        #且[]中出现了数字,他就是降维操作,所以output由三维降到两维。
        logits_clsf=self.classifier(h_pooled)
        #得到判断两句话是不是上下句关系的结果

        #得到被mask位置的词,准备与正确词进行比较
        masked_pos=masked_pos[:,:,None].expand(-1,-1,d_model) #[batch_size,max_pred,d_model]
        h_masked=torch.gather(output,1,masked_pos) #masking position [batch_size,max_pred,d_model]
        h_masked=self.activ2(self.linear(h_masked))# [batch_size,max_pred,d_model]
        logits_lm=self.fc2(h_masked) #[batch_size,max_pred,vocab_size]
        return logits_lm,logits_clsf #预测的被mask地方的值,预测两句话是否为上下句的结果

model=BERT()
criterion=torch.nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(model.parameters(),lr=0.001)

model.to(device)

for epoch in range(50):
    for input_ids,segment_ids,masked_tokens,masked_pos,isNext in trainloader:
        input_ids=input_ids.to(device)
        segment_ids=segment_ids.to(device)
        masked_pos=masked_pos.to(device)
        masked_tokens=masked_tokens.to(device)
        isNext=isNext.to(device)
        logits_lm,logits_clsf=model(input_ids,segment_ids,masked_pos)
        loss_lm=criterion(logits_lm.view(-1,vocab_size),masked_tokens.view(-1))
        loss_lm=(loss_lm.float())/5
        loss_clsf=criterion(logits_clsf,isNext)
        loss=loss_clsf+loss_lm
        if (epoch+1)%10==0:
            print('Epoch: %04d'%(epoch+1),'loss=','{:.6f}'.format(loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

input_ids,segment_ids,masked_tokens,masked_pos,isNext=batch[1]
input_word=[idx2word[i] for i in input_ids]
input_ids=torch.LongTensor([input_ids]).to(device)
segment_ids=torch.LongTensor([segment_ids]).to(device)
masked_pos=torch.LongTensor([masked_pos]).to(device)
masked_tokens=torch.LongTensor(masked_tokens).to(device)
isNext=torch.LongTensor([isNext]).to(device)
print(text)
print('================================')
print([idx2word[w.item()] for w in input_ids[0] if idx2word[w.item()]!='[PAD]'])

logits_lm,logits_clsf=model(input_ids,segment_ids,masked_pos)
logits_lm=logits_lm.data.max(2)[1][0].data
#最大值总是在第0个,代码应该是有问题的。且5个mask的值完全一样,所以明天好好看看代码
date=[i.item() for i in logits_lm]
for i in date:
    if i:
        print(f'被mask掉的词是{idx2word[i]}')
"""
max(2):再第2维上取最大值,取完后维度:[batch,max_pred],[batch,max_pred]。此时有两个维度值。第一个是具体的值,第二个是位置。(值不需要,就是一个经过了softmax后比其他值大的数)
[1]:取到最大值对应的位置。维度是:[batch,max_pred]
[0]:因为batch为1(只取了第一组值),所以此时维度是:[max_pred]
注:max会使tensor降维
"""
print('masked token list:',[pos.item() for pos in masked_tokens if pos!=0])
print('predict masked tokens list:',[pos.item() for pos in logits_lm if pos!=0])

logits_clsf=logits_clsf.data.max(1)[1].data[0]
print('isNext:',True if isNext.item() else False)
print('predict isNext:',True if logits_clsf else False)

Original: https://blog.csdn.net/buster120/article/details/116057495
Author: buster120
Title: 在pytorch上实现bert的简单预训练过程

原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/547989/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球