Transformer模型实现翻译任务

Transformer简明教程, 从理论到代码实现到项目实战, NLP进阶必知必会.

https://www.bilibili.com/list/ml3086466753?oid=253396522&bvid=BV19Y411b7qx

程序包括data.py、mask.py、util.py、model.py、main.py

1、data.py 生成预测数据

#定义试验数据
#定义模型
import random
import numpy as np
import torch

#定义字典 class "str"
zidian_x="<SOS>,<EOS>,<PAD>,0,1,2,3,4,5,6,7,8,9,q,w,e,r,t,y,u,i,o,p,a,s,d,f,g,h,j,k,l,z,x,c,v,b,n,m"


# calss "dict"
zidian_x = {word: i for i, word in enumerate(zidian_x.split(","))}
zidian_xr=[k for k, v in zidian_x.items()]

zidian_y = {k.upper():v for k,v in zidian_x.items()}
zidian_yr=[k for k,v in zidian_y.items()]


#生成数据的函数
def get_data():
    #定义词集合
    words=[
        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'q', 'w', 'e', 'r',
        't', 'y', 'u', 'i', 'o', 'p', 'a', 's', 'd', 'f', 'g', 'h', 'j', 'k',
        'l', 'z', 'x', 'c', 'v', 'b', 'n', 'm'
    ]

    #定义每个词被选中的概率
    p=np.array([
        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
    ])

    #计算概率
    p=p/p.sum()

    #随机选n个词
    n=random.randint(30,48)
    x=np.random.choice(words,size=n,replace=True,p=p)

    #采样结果就是x 将数组或矩阵对象转换为列表
    x=x.tolist()

    #y是对x的变换得到的
    #字母大写，数据取10以内的互补数
    def f(i):
        i=i.upper()
        if not i.isdigit():
            return i;

        i=9-int(i)
        return  str(i)

    y=[f(i) for i in x]

    y=y+[y[-1]]

    #逆序
    y=y[::-1]

    #加上首尾符号
    x = ['<SOS>'] + x + ['<EOS>']
    y = ['<SOS>'] + y + ['<EOS>']

    #补充pad道固定长度

    x = x + ['<PAD>'] * 50
    y = y + ['<PAD>'] * 51
    x=x[:50]
    y=y[:51]

    #编码成数据
    x = [zidian_x[i] for i in x]
    y = [zidian_y[i] for i in y]

    #转换成tensor
    x = torch.LongTensor(x)
    y = torch.LongTensor(y)

    return x,y

#定义数据集
class Dataset(torch.utils.data.Dataset):
    def __init__(self):
        super(Dataset,self).__init__()

    def __len__(self):
        return 100000

    def __getitem__(self, i):
        return get_data()

loader=torch.utils.data.DataLoader(Dataset(),
                                   batch_size=8,
                                   drop_last=True,
                                   shuffle=True,
                                   collate_fn=None)

2、mask.py 生成遮挡矩阵

#定义mask函数
import torch

from data import zidian_x,zidian_y
def mask_pad(data):
    #b句话，每句话50各词，这里还没有embed
    #data=[b,50]
    #判断每个词是不是<PAD>
    mask = data == zidian_x['<PAD>']

    #[b,50]=>[b,1,1,50]
    mask=mask.reshape(-1,1,1,50)

    #在计算注意力时，是计算的50个词和50个词相互之间的注意力，所以是50X50的矩阵
    #是PAD的列是True ，意味着任何词对PAD的注意力都是0
    #但是PAD本身对其他的词的注意力并不是0
    #所以是PAD的行不是0

    #复制n次
    #[b, 1, 1, 50]=>[b, 1, 50, 50]=>

    mask=mask.expand(-1,1,50,50)

    return mask

def mask_tril(data):
    #b句话,每句话50个词，这里面还没有embed的
    #data=[b,50]
    #50X50矩阵表示每个词对其他词是否可见
    #上三角矩阵，不包括对角线，意味着对每个词而言，他只能看到他自己，和他之前的词，看不到之后的词
    #[1, 50, 50]

    tril=1-torch.tril(torch.ones(1,50,50,dtype=torch.long))

    #判断当Y中每个词是不是PAD，如果是Pad则不可见
    #[b,50]
    mask = data == zidian_y['<PAD>']

    #变形 为了之后的计算
    #[b,1,50]
    mask=mask.unsqueeze(1).long()

    #mask和tril求并集
    # [b,1,50]+[1,50,50]->[b,50,50]
    mask=mask+tril

    #转布尔型
    mask=mask>0

    #换布尔型 ，增加一个维度，便于后续计算
    mask=(mask==1).unsqueeze(dim=1)

    return mask

3、util.py 包括位置编码和多头注意力

#定义工具函数
import math

import torch
#注意力计算函数
def attention(Q,K,V,mask):
    #b句话，每句话50个词，每个词编码成32维向量，4个头，每个头分到8维向量
    #Q,K,V=[b,4,50,8]

    #[b,4,50,8]*[b,4,8,50]=>[b,4,50,50]
    #permute 是转置函数
    #Q,K矩阵相乘，求每个词相对其他所有词的注意力
    score=torch.matmul(Q,K.permute(0,1,3,2))

    #除以每个头尾数的平方根，做数值缩放
    score/=8**0.5

    #mask遮盖，mask是true的地方都被替换成-inf,这样在计算softmax时候，-inf会被压缩到0
    #mask=[b,1,50,50]

    score=score.masked_fill_(mask,-float('inf'))
    #经过sofmax函数
    score=torch.softmax(score,dim=-1)

    #以注意力分数乘以V，得到最终的注意力结果
    #[b,4,50,50]*[b,4,50,8]=>[b,4,50,8]
    score=torch.matmul(score,V)

    #每个头的计算结果合并到一起
    #[b,4,50,8]=>[b,50,32]

    score=score.permute(0,2,1,3).reshape(-1,50,32)

    return score

    #LayerNorm用在自然语言中
    #BatchNorm用在图像处理中

#多头注意力计算层
class MultiHead(torch.nn.Module):
    def __init__(self):
        super().__init__()
        #Q的全连接层
        self.fc_Q = torch.nn.Linear(32, 32)
        #K的全连接层
        self.fc_K = torch.nn.Linear(32, 32)
        #V的全连接层
        self.fc_V = torch.nn.Linear(32, 32)

        #输出层
        self.out_fc = torch.nn.Linear(32, 32)

        #归一化层
        self.norm=torch.nn.LayerNorm(normalized_shape=32,elementwise_affine=True)

        #dropout层 防止过拟合
        self.dropout=torch.nn.Dropout(p=0.1)

    def forward(self,Q,K,V,mask):
        #b句话，每句话50个词，每个词编码成32维向量
        #Q,K,V=[b,50,32]
        b = Q.shape[0]

        #保留下原始的Q,后面要用来短接
        clone_Q = Q.clone()
        #规范化
        Q = self.norm(Q)
        K = self.norm(K)
        V = self.norm(V)


        #线性运算，维度不变
        #[b, 50, 32]=>[b, 50, 32]
        Q = self.fc_Q(Q)
        K = self.fc_K(K)
        V = self.fc_V(V)

        #拆分成多个头
        #b句话，每句话50个词，每个词编码成32维向量，4个头，每个头分到8维向量
        #[b,50,32]=>[b, 4, 50, 8]
        Q = Q.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)
        K = K.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)
        V = V.reshape(b, 50, 4, 8).permute(0, 2, 1, 3)

        #计算注意力
        #[b, 4, 50, 8]=>[b,50,32]
        score=attention(Q,K,V,mask)


        #计算输出，维度不变
        # [b, 50, 32]=>[b, 50, 32]
        score=self.dropout(self.out_fc(score))


        #短接
        score=clone_Q+score
        return score

#位置编码层
class PositionEmbedding(torch.nn.Module):
    def __init__(self):
        super().__init__()

        #pos是第几个词，i是第几个维度，d_model是维度总数

        def get_pe(pos,i,d_model):
            fenmu=1e4**(i/d_model)
            pe=pos/fenmu

            if i%2==0:
                return math.sin(pe)
            return math.cos(pe)

        #初始化位置编码矩阵
        pe=torch.empty(50,32)
        for i in range(50):
            for j in range(32):
                pe[i,j]=get_pe(i,j,32)
        pe=pe.unsqueeze(0)

        #定义为不更新常量
        self.register_buffer('pe',pe)

        #词编码层
        self.embed=torch.nn.Embedding(39,32)
        #初始化参数
        self.embed.weight.data.normal_(0,0.1)


    def forward(self,x):
        #8句话 每句话50个词
        #[8,50]=>[8,50,32]
        embed=self.embed(x)

        #词编码和位置编码相加
        #[8,50,32]+[1,50,32]->[8,50,32]
        embed=embed+self.pe
        return embed

class FullConnectedOutput(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc=torch.nn.Sequential(
            torch.nn.Linear(in_features=32, out_features=64),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=64, out_features=32),
            torch.nn.Dropout(p=0.1)
        )

        self.norm=torch.nn.LayerNorm(normalized_shape=32,
                                     elementwise_affine=True)
    def forward(self,x):
        clone_x=x.clone()
        #保留原始的X，后面用做短接

        #规范化
        x=self.norm(x)

        #输入全连接 形状不变
        #[b,50,32]=>[b,50,32]
        out=self.fc(x)

        #做短接
        out=clone_x+out

        return out

4、model.py 定义编码器、解码器和模型

import torch.nn
from util import MultiHead,FullConnectedOutput,PositionEmbedding
from mask import mask_pad,mask_tril
#编码器层
class EncoderLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mh=MultiHead()
        self.fc=FullConnectedOutput()
    def forward(self,x,mask):
        #计算自注意力，维度不变
        #[b,50,32]=>[b,50,32]
        score=self.mh(x,x,x,mask)

        #全连接输出，维度不变
        # [b,50,32]=>[b,50,32]
        out=self.fc(score)

        return out

class Encoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_1=EncoderLayer()
        self.layer_2=EncoderLayer()
        self.layer_3=EncoderLayer()
    def forward(self,x,mask):
        x=self.layer_1(x,mask)
        x=self.layer_2(x,mask)
        x=self.layer_3(x,mask)
        return x
#解码器层
class DecoderLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.mh1 = MultiHead()
        self.mh2 = MultiHead()
        self.fc = FullConnectedOutput()
    def forward(self,x,y,mask_pad_x,mask_tril_y):
        #先计算y的自注意力，维度不变
        # [b,50,32]=>[b,50,32]
        y=self.mh1(y,y,y,mask_tril_y)

        #结合x和y的注意力计算，维度不变
        # [b,50,32],[b,50,32]=>[b,50,32]
        y = self.mh1(y, x, x, mask_pad_x)

        #全连接输出，维度不变
        # [b,50,32]=>[b,50,32]
        y=self.fc(y)

        return y

class Decoder(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer_1=DecoderLayer()
        self.layer_2=DecoderLayer()
        self.layer_3=DecoderLayer()
    def forward(self,x,y,mask_pad_x,mask_tril_y):
        y=self.layer_1(x,y,mask_pad_x,mask_tril_y)
        y=self.layer_2(x,y,mask_pad_x,mask_tril_y)
        y=self.layer_3(x,y,mask_pad_x,mask_tril_y)

        return y


#主模型
class Transformer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.embed=PositionEmbedding()
        self.encoder=Encoder()
        self.decoder=Decoder()
        self.fc_out=torch.nn.Linear(32,39)
    def forward(self,x,y):
        #mask形状[b,1,50,50]
        mask_pad_x=mask_pad(x)
        mask_tril_y=mask_tril(y)

        #对数据编码，添加位置信息
        #x=[b,50]=>[b,50,32]
        x=self.embed(x)
        y=self.embed(y)


        #编码计算
        # [b,50,32]=>[b,50,32]
        x=self.encoder(x,mask_pad_x)

        #解码层计算
        # [b,50,32],[b,50,32]=>[b,50,32]
        y=self.decoder(x,y,mask_pad_x,mask_tril_y)


        #全连接输出
        # [b,50,32]=>[b,50,39]
        y=self.fc_out(y)

        return y

5、main.py 包括训练函数和预测函数

#主程序 试验+测试
import torch
from model import Transformer
from data import loader,zidian_y,zidian_xr,zidian_yr
from mask import mask_pad,mask_tril
model=Transformer()
loss_func=torch.nn.CrossEntropyLoss()
optim=torch.optim.Adam(model.parameters(),lr=2e-3)
sched=torch.optim.lr_scheduler.StepLR(optim,step_size=3,gamma=0.5)

for epoch in range(1):
    for i,(x,y) in enumerate(loader):
        #x=[8,50]
        #y=[8,51]

        #在训练时，是拿y的每一个字符输入，预测下一个字符，因此不需要最后一个字符
        pred=model(x,y[:,:-1])
        #得到数据形状[8,50,39]

        # [8,50,39]=>[400,39]

        pred=pred.reshape(-1,39)

        #[8,51]=>[400]
        y=y[:,1:].reshape(-1)

        #忽略pad
        select=y!=zidian_y['<PAD>']

        pred=pred[select]
        y=y[select]

        loss=loss_func(pred,y)

        optim.zero_grad()

        loss.backward()
        optim.step()

        if i%200==0:
            # [select,39]=>[select]
            pred=pred.argmax(1)
            correct=(pred==y).sum().item()
            accuracy=correct/len(pred)

            lr=optim.param_groups[0]['lr']
            print(epoch,i,lr,loss.item(),accuracy)


    sched.step()

#预测函数
def predict(x):
    # x=[1,50] 模型评估模式
    model.eval()
    mask_pad_x=mask_pad(x)

    #初始化输出，这个是固定值
    #[1,50]
    #[[0,2,2,2,2,2....]]

    target=[zidian_y['<SOS>']]+[zidian_y['<PAD>']]*49
    target=torch.LongTensor(target).unsqueeze(0)

    #x编码，添加位置信息
    #[1,50]=>[1,50,32]
    x=model.embed(x)

    #编码层计算，维度不变
    #[1,50,32]->[1,50,32]

    x=model.encoder(x,mask_pad_x)

    for i in range(49):

        #[1,50]
        y=target

        #[1,1,50,50]
        mask_tril_y=mask_tril(y)

        #y编码 添加位置新
        #[1,50]=>[1,50,32]
        y=model.embed(y)

        #解码层计算，维度不变
        #[1,50,32],[1,50,32]=>[1,50,32]

        y=model.decoder(x,y,mask_pad_x,mask_tril_y)

        #全连接输出 39分类
        #[1,50,32]=>[1,50,39]
        out=model.fc_out(y)

        #取出当前词的输出
        out=out[:,i,:]


        #取出分类结果
        #[1:39]=>[1]
        out=out.argmax(dim=1).detach()

        #以当前次预测下一个词，填充道结果中
        target[:,i+1]=out

    return target


for i,(x,y) in enumerate(loader):
    break

for i in range(8):
    print(i)
    print(''.join([zidian_xr[i] for i in x[i].tolist()]))
    print(''.join([zidian_yr[i] for i in y[i].tolist()]))
    print(''.join([zidian_yr[i] for i in predict(x[i].unsqueeze(0))[0].tolist()]))