Transformer Model Architecture代码

本文所呈现代码，仅为复现模型各部分代码，不用于训练。如果需要进行训练的完整代码，请见代码块中的注释部分。

#完整可运行代码，详见https://github.com/hyunwoongko/transformer
import numpy as np
import torch
import torch.nn as nn
import math
#Positional Encoding
class PositionalEncoding(nn.Module):
    def __init__(self,max_len,d_model,device):
        """
        加载位置编码,用于宇input embedding相加
        :param max_len: 序列的最大长度，也就是最大的位置编码
        :param d_model:模型的dimension,论文里是512
        :param device:cpu or cuda
        """
        super(PositionalEncoding,self).__init__()
        self.encoding=torch.zeros(max_len,d_model,device=device)
        self.encoding.requires_grad=False
        #位置编码不需要计算梯度
        pos=torch.arange(0,max_len,device=device)#生成序列，左闭右开
        pos=torch.unsqueeze(pos,1).to(torch.float32)#转换成2维，变成浮点数

        _2i=torch.arange(0,d_model,device=device)
        _2i=torch.unsqueeze(pos,1).to(torch.float32)

        self.encoding[:,0::2]=torch.sin(pos/(10000**(_2i/d_model)))#偶数列赋值
        self.encoding[:,1::2]=torch.cos(pos/(10000**(_2i/d_model)))#奇数列赋值

    def forward(self,x):
        batch_size,seq_len=x.size()
        #seq_len是实际输入序列长度小于等于max_len
        return self.encoding[:seq_len,:]

#Scaled Dot-Product Attention
class ScaledDotProductAttention(nn.Module):
    """
    计算注意力分数
    """
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()
        self.softmax=nn.Softmax(dim=-1)

    def forward(self,q,k,v,mask=None,e=1e-12):
        #由于在multi-head attention开头，对q\k\v都进行了增维处理，也就是有多个head
        #[batch,head,seq_len,d_k/v/q],d_model
        batch,head,seq,d_k=k.size
        k_t=k.transpose(2,3)
        score=(q@k_t)/math.sqrt(d_k)#@表示点积，得分矩阵的形状是seq_len*seq_len
        #对于decoder部分，还会有mask
        if mask is not None:
            score=score.masked_fill(mask==0,-10000)#给得分矩阵为0的地方赋值为极小值
        score=self.softmax(score)
        v=score@v
        return v,score

#Multi-Head Attention
class MultiHeadAttention(nn.Module):
    def __init__(self,d_model,d_head):
        super(MultiHeadAttention, self).__init__()
        self.head=d_head
        self.w_q=nn.Linear(d_model,d_model)
        self.w_k=nn.Linear(d_model,d_model)
        self.w_v=nn.Linear(d_model,d_model)
        self.attention=ScaledDotProductAttention()
        self.w_cat=nn.Linear(d_model,d_model)

    def forward(self,q,k,v,mask=None):
        q,k,v=self.w_q(q),self.w_k(k),self.w_v(v)
        #把q,k,v拆分成n个head的形式
        q,k,v=self.split(q),self.split(k),self.split(v)
        value,score=self.attention(q,k,v,mask=mask)
        value=self.concat(value)
        out=self.w_cat(value)
        return out
    def split(self,tensor):
        """
        切分向量
        :param tensor: [batch,seq_len,d_model]
        :return: [batch,head,seq_len,d_tensor],d_tensor=d_q/k/v
        """
        batch_size,seq_len,d_model=tensor.size()
        d_tensor=d_model//self.head
        tensor=tensor.view(batch_size,seq_len,self.head,d_tensor).transpose(1,2)
        return tensor
    def concat(self,tensor):
        """
        split的反向操作
        :param tensor: [batch,head,seq_len,d_tensor]
        :return: [batch,seq_len,d_model]
        """
        batch_size,head,seq_len,d_tensor=tensor.size()
        d_model=head*d_tensor
        tensor=tensor.transpose(1,2).contiguous().view(batch_size,seq_len,d_model)#contiguous() 确保了内存的连续性
        return tensor

class LayerNorm(nn.Module):
    def __init__(self,d_model,eps=1e-12):
        """
        对每个样本的特征维求normalization
        batch_norm是对每个minibatch做normalization
        :param d_model: 特征维上的数量
        :param eps:一个小常数，避免分母除0错误
        """
        super(LayerNorm, self).__init__()
        self.eps=eps
        self.gamma=nn.Parameter(torch.ones(d_model))#将张量包装成一个模型的参数。这使得张量在模型训练过程中可以被优化器更新。
        self.beta=nn.Parameter(torch.zeros(d_model))
    def forward(self,x):
        #x shape[batch_size,seq_length,d_model]
        #对样本的特征维（隐藏层那一个维度）求均值和方差
        mean=x.mean(-1,keepdim=True)#-1表示最后一个维度，keepdim表示结果是否与原维度相同
        #mean的shape[batch_size,seq_length,1]
        var=x.var(-1,unbiased=True,keepdim=True)#unbiased表示无偏估计
        out=(x-mean)/np.sqrt(var+self.eps)
        out=self.gamma*out+self.beta
        return out

class PointwiseFeedForward(nn.Module):
    def __init__(self,d_model,d_ff=2048,drop_rate=0.1):
        """
        2个全连接层中间跟一个relu激活函数
        :param d_model: input 和output
        :param d_ff: 中间层
        :param drop_rate: dropout的概率
        """
        super(PointwiseFeedForward, self).__init__()
        self.layer1=nn.Linear(d_model,d_ff)
        self.layer2=nn.Linear(d_ff,d_model)
        self.relu=nn.ReLU()
        self.dropout=nn.Dropout(p=drop_rate)

    def forward(self,x):
        x=self.layer1(x)
        x=self.relu(x)
        x=self.dropout(x)
        x=self.layer2(x)
        return x

#encoder块
class EncoderLayer(nn.Module):
    def __init__(self,d_model,n_heads,d_ff,drop_rate):
        super(EncoderLayer, self).__init__()
        self.attention=MultiHeadAttention(d_model=d_model,d_head=n_heads)
        self.norm1=LayerNorm(d_model=d_model)
        self.dropout1=nn.Dropout(p=drop_rate)

        self.ff=PointwiseFeedForward(d_model=d_model,d_ff=d_ff,drop_rate=drop_rate)
        self.norm2=LayerNorm(d_model=d_model)
        self.dropout2=nn.Dropout(p=drop_rate)

    def forward(self,x,src_mask):
        _x=x
        x=self.attention(q=x,k=x,v=x,mask=src_mask)
        x=self.dropout1(x)
        x=self.norm1(_x+x)
        _x=x
        x=self.ff(x)
        x=self.norm2(_x+x)
        return x
#Encoder
class Encoder(nn.Module):
    def __init__(self,enc_voc_size,max_len,d_ff,n_head,d_model,drop,device,n_layers=6):
        super(Encoder, self).__init__()
        #embedding操作
        #生成encoder
        self.layers=nn.ModuleList([EncoderLayer(d_model=d_model,d_ff=d_ff,n_heads=n_head,drop_rate=drop) for _ in range(n_layers)])

    def forward(self,x,src_mask):
        #进行embedding
        #放入每个encoder层
        for layer in self.layers:
            x=layer(x,src_mask)

        return x

#decoder layer
class DecoderLayer(nn.Module):
    def __init__(self,d_model,d_ff,n_head,drop_rate):
        super(DecoderLayer, self).__init__()
        self.attention=MultiHeadAttention(d_model=d_model,d_head=n_head)
        self.norm1=LayerNorm(d_model=d_model)
        self.dropout1=nn.Dropout(p=drop_rate)

        self.enc_dec_attention=MultiHeadAttention(d_model=d_model,d_head=n_head)
        self.norm2=LayerNorm(d_model=d_model)
        self.dropout2=nn.Dropout(p=drop_rate)

        self.ff=PointwiseFeedForward(d_model=d_model,d_ff=d_ff,drop_rate=drop_rate)
        self.norm3=LayerNorm(d_model=d_model)
        self.dropout3=nn.Dropout(p=drop_rate)

    def forward(self,dec,enc,tgt_mask,src_mask):
        _x=dec
        x=self.attention(q=dec,k=dec,v=dec,mask=tgt_mask)
        x = self.dropout1(x)
        x=self.norm1(_x+x)

        if enc is not None:
            _x=x
            x=self.enc_dec_attention(q=x,k=enc,v=enc,mask=src_mask)

            x=self.dropout2(x)
            x=self.norm2(_x+x)
        _x=x
        x=self.ff(x)
        x=self.dropout3(x)
        x=self.norm3(_x+x)
        return x

#decoder
class Decoder(nn.Module):
    def __init__(self,dec_voc_size,max_len,d_model,d_ff,n_head,drop_rate,n_layers):
        super(Decoder, self).__init__()
        #embedding
        #生成decoder layer
        self.layer=nn.ModuleList([DecoderLayer(d_model=d_model,d_ff=d_ff,n_head=n_head,drop_rate=drop_rate) for _ in range(n_layers)])
        self.linear=nn.Linear(d_model,dec_voc_size)

    def forward(self,tgt,src,tgt_mask,src_mask):
        #对tgt embedding
        for layers in self.layer:
            tgt=layers(tgt,src,tgt_mask,src_mask)

        output=self.linear(tgt)
        output=nn.Softmax(output)
        return output