CodeByMySelf-Transformer

所有之前：CodeByMySelf 系列之 Transformer 的实现记录，只有标题没有内容的是TODO，针对经典算法的变体和改进算法在其他文章。

Transformer

下面是一个最基本的Transformer实现，主要包含以下几个部分：

基本模块
- 多头注意力模块
- Token和位置嵌入模块
- LayerNorm
- 前馈神经网络FFN
$EncoderLayer -> Encoder$
$DecoderLayer -> Decoder$
$Encoder+Decoder -> Transformer$

多头注意力MHA

多头注意力机制的公式如下：

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O $$

其中，每一个头的计算分别为：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V $$

对于图像数据和序列数据，多头注意力机制的实现是不同，但是核心原理是相同的。对于序列数据，每个位置对应一个token；对于图像数据，每个位置对应一个patch块或者一个像素块，如果是$patch$块，需要展开到一维向量。这里实现的为序列数据的多头注意力机制,初始$shape$为$(batch,time,d_model)$


import torch
from torch import nn
import torch.functional as F
import math

class MultiHeadAttention(nn.Module):
    def __init__(self,d_model,n_head):
        super(MultiHeadAttention, self).__init__()
        
        assert d_model % n_head == 0
        self.d_model = d_model
        self.n_head = n_head

        self.w_k = nn.Linear(d_model,d_model)
        self.w_q = nn.Linear(d_model,d_model)
        self.w_v = nn.Linear(d_model,d_model)
        self.w_conbine = nn.Linear(d_model,d_model,bias=False) # 对应W^O
        self.softmax = nn.Softmax(dim=-1)

    def forward(self,q,k,v,mask=None):
        batch,time,dimension = q.shape() # q,k,v shape: (batch,time,d_model)
        n_d = self.d_model // self.n_head

        q,k,v = self.w_q(q),self.w_k(k),self.w_v(v)
        print(q.shape, batch, time, self.n_head, n_d)
        q = q.view(batch,time,self.n_head,n_d).transpose(1,2) # self.n_head * n_d = d_model
        k = k.view(batch,time,self.n_head,n_d).transpose(1,2) # shape: (batch,n_head,time,n_d)
        v = v.view(batch,time,self.n_head,n_d).transpose(1,2)

        # compute attention by the formula
        score = q @ k.transpose(-2,-1) / math.sqrt(n_d) # (batch,n_head,time,time)
        if mask is not None:
            score = score.masked_fill(mask==0,float('-inf'))    
        attn = self.softmax(score) @ v # (batch,n_head,time,n_d)
        attn = attn.transpose(1,2).contiguous().view(batch,time,dimension) # (batch,time,dimension)
        attn = self.w_conbine(attn) # (batch,time,dimension)
        return attn

d_model = 512
n_head = 8
mha = MultiHeadAttention(d_model,n_head)
x = torch.randn(2,10,d_model)
y = mha(x,x,x) # self-multi-head-attention
print(f"shape: {y.shape}\n输出:{y}") # torch.Size([2, 10, 512])

Token和位置嵌入模块

Token Embedding 是将离散的词（token）转换为连续的向量表示的过程。通常通过查找表（如 nn.Embedding）将每个词的索引映射为一个高维向量，使模型能够处理和学习词之间的语义关系。公式表示：

$$ \text{TokenEmbedding}(x) = E[x] $$

其中 $E$ 是嵌入矩阵，$x$ 是词的索引。TokenEmbedding是一个专门的领域，这里不做深究，直接使用torch自己的类。

Position Embedding为模型提供序列的顺序信息，Transformer 结构本身不具备序列顺序信息，因此需要引入 Position Embedding，将每个位置的信息编码为向量，并与 Token Embedding 相加。常见做法有两种：

可学习的位置嵌入（Learnable Position Embedding）：每个位置对应一个可训练的向量。
固定的位置嵌入（Sinusoidal Position Embedding）：使用三角函数编码序列位置信息。

以固定的位置嵌入为例,公式如下:

$$ PE_{(pos,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$$$ PE_{(pos,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$

import torch
from torch import nn
import torch.functional as F
import math

class TokenEmbedding(nn.Embedding):
    def __init__(self,vocab_size,d_model):
        super(TokenEmbedding,self).__init__(vocab_size,d_model,padding_idx=1)
#----------------------------------------------------------------------------
class PositionalEncoding(nn.Module):
    def __init__(self,d_model,device,max_len=5000):
        super(PositionalEncoding,self).__init__()

        self.encoding = torch.zeros(max_len,d_model,device=device) # (max_len,d_model)
        self.encoding.requires_grad = False # 固定的位置嵌入不需要梯度

        pos = torch.arange(0,max_len,device=device).float().unsqueeze(1) # (max_len)-> (max_len,1)
        _2i = torch.arange(0,d_model,2,device=device).float() # (d_model/2,)

        # PE(pos,2i) = sin(pos/10000^(2i/d_model))
        self.encoding[:,0::2] = torch.sin(pos / (10000 ** (_2i/d_model)))
        # PE(pos,2i+1) = cos(pos/10000^(2i/d_model))
        self.encoding[:,1::2] = torch.cos(pos / (10000 ** (_2i/d_model)))
    
    def forward(self,x):
        batch,seq_len,dimension = x.size()
        # return x + self.encoding[:seq_len,:] # 直接返回带位置嵌入的x,多余的位置嵌入不需要
        return self.encoding[:seq_len,:] # 只返回位置嵌入
#----------------------------------------------------------------------------
class TransformerEmbedding(nn.Module):
    def __init__(self,vocab_size,d_model,device,max_len=5000,dropout_p=0.1):
        super(TransformerEmbedding,self).__init__()
        self.token_embedding = TokenEmbedding(vocab_size,d_model)
        self.position_embedding = PositionalEncoding(d_model,device,max_len)
        self.dropout = nn.Dropout(dropout_p)

    # def forward(self,x):
    #     x = self.token_embedding(x)
    #     x = self.position_embedding(x)
    #     x = self.dropout(x)
    #     return x
    
    def forward(self,x):
        t_emb = self.token_embedding(x)
        p_emb = self.position_embedding(t_emb)
        return self.dropout(t_emb + p_emb)

LayerNorm

LayerNorm（层归一化）是一种常用的归一化方法，主要用于深度学习模型中提升训练稳定性和加速收敛。它的作用是对每个样本的特征维度进行归一化，使其均值为0、方差为1，并引入可学习的缩放和偏移参数。

公式如下：

$$ \text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta $$

$x$：输入向量 $\mu$：均值，$\mu = \frac{1}{N} \sum_{i=1}^N x_i$ $\sigma^2$：方差，$\sigma^2 = \frac{1}{N} \sum_{i=1}^N (x_i - \mu)^2$ $\gamma$ 和 $\beta$：可学习的缩放和偏移参数 $\epsilon$：防止除零的小常数

Tip:LayerNorm（层归一化）和 BatchNorm（批归一化）主要区别如下： BatchNorm 是对每个特征维度在整个 batch 上做归一化，依赖于 batch 的统计信息。 LayerNorm 是对每个样本的所有特征维度做归一化，和 batch 大小无关，对于序列数据(batch,time,d)来说，就是对每个time的d维数据进行归一化；对于图像数据来说，就是在图像每个通道的像素点进行归一化。

import torch
from torch import nn
import torch.functional as F
import math

class LayerNorm(nn.Module):
    def __init__(self,d_model,eps=1e-9):
        super(LayerNorm,self).__init__()

        # two learnable parameters
        self.gamma = nn.Parameter(torch.ones(d_model))
        self.beta = nn.Parameter(torch.zeros(d_model))
        
        # one small constant to prevent division by zero
        self.eps = eps

    def forward(self,x):
        mean = x.mean(-1,keepdim=True) # (batch,time,1)
        std = x.std(-1,keepdim=True) # (batch,time,1)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta # by the formula

x = torch.randn(2,10,512)
ln = LayerNorm(512)
y = ln(x)
print(f"shape: {y.shape}\n输出:{y}")

FFN前馈网络（别看名字，实际很简单）

FFN 的结构如下：

包含两个线性变换（全连接层）和一个激活函数（通常是ReLU或GELU）。
对每个时间步/位置单独处理，不进行序列间的信息交互。

公式表示：

$$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$

import torch
from torch import nn
import torch.functional as F
import math

class ffn(nn.Module):
    def __init__(self,d_model,d_hidden,dropout=0.1):
        super(ffn,self).__init__()

        self.fc1 = nn.Linear(d_model,d_hidden)
        self.fc2 = nn.Linear(d_hidden,d_model)
        self.dropout = nn.Dropout(dropout) # dropout放在激活函数后，减少激活的神经元
        self.relu = nn.ReLU()

    def forward(self,x):
        return self.fc2(self.dropout(self.relu(self.fc1(x)))) # by the formula

集成所有模块，构建完整的Transformer模型

encoder结构

decoder结构

首先构建EncoderLayer和DecoderLayer，然后将它们堆叠成Encoder和Decoder。

import torch
from torch import nn
import torch.functional as F
import math

class EncoderLayer(nn.Module):
    def __init__(self,d_model,n_head,d_hidden,dropout_p=0.1):
        super(EncoderLayer,self).__init__()
        
        self.mha = MultiHeadAttention(d_model,n_head)
        self.ln1 = LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout_p)

        self.ffn = ffn(d_model,d_hidden,dropout_p)
        self.ln2 = LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout_p)

    def forward(self,x,mask=None):
        _x = x
        x = self.mha(x,x,x,mask)
        x = self.dropout(x)
        x = self.ln1(x + _x)

        _x = x
        x = self.ffn(x)
        x = self.dropout2(x)
        x = self.ln2(x + _x)

        return x

class DecoderLayer(nn.Module):
    def __init__(self,d_model,n_head,d_hidden,dropout_p=0.1):
        super(DecoderLayer,self).__init__()

        self.s_mha = MultiHeadAttention(d_model,n_head) # self-multi-head-attention with mask
        self.ln1 = LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout_p)

        self.cross_mha = MultiHeadAttention(d_model,n_head) # cross-multi-head-attention
        self.ln2 = LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout_p)

        self.ffn = ffn(d_model,d_hidden,dropout_p)
        self.ln3 = LayerNorm(d_model)
        self.dropout3 = nn.Dropout(dropout_p)

    def forward(self,dec,enc,time_mask=None,src_mask=None):
        _x = dec
        x = self.s_mha(dec,dec,dec,time_mask) # 避免解码器看到未来的信息,下三角掩码

        x = self.dropout1(x)
        x = self.ln1(x + _x)

        if enc is not None:
            _x = x
            x = self.cross_mha(x,enc,enc,src_mask) # 对于padding的mask

            x = self.dropout2(x)
            x = self.ln2(x + _x)
        
        _x = x
        x = self.ffn(x)

        x = self.dropout3(x)
        x = self.ln3(x + _x)
        return x

首先使用之前写好的组件构成Encoder和Decoder,然后完成整个Transformer，Mask是Transformer中最关键的一部分,怎么构建两个掩码以及两个掩码的作用分别是什么

class Encoder(nn.Module):
    def __init__(self,vocab_size,d_model,n_head,d_hidden,device,max_len=5000,dropout_p=0.1,num_layers=6):
        super(Encoder,self).__init__()

        self.embedding = TransformerEmbedding(vocab_size,d_model,device,max_len,dropout_p)
        self.layers = nn.ModuleList(
            [EncoderLayer(d_model,n_head,d_hidden,dropout_p) for _ in range(num_layers)]
        )

    def forward(self,x,src_mask=None):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x,src_mask)
        return x

class Decoder(nn.Module):
    def __init__(self,vocab_size,d_model,n_head,d_hidden,device,max_len=5000,dropout_p=0.1,num_layers=6):
        super(Decoder,self).__init__()

        self.embedding = TransformerEmbedding(vocab_size,d_model,device,max_len,dropout_p)
        self.layers = nn.ModuleList(
            [DecoderLayer(d_model,n_head,d_hidden,dropout_p) for _ in range(num_layers)]
        )
        self.fc_out = nn.Linear(d_model,vocab_size,bias=False)

    def forward(self,dec,enc,time_mask=None,src_mask=None):
        x = self.embedding(dec)
        for layer in self.layers:
            x = layer(x,enc,time_mask,src_mask)
        x = self.fc_out(x)
        return x

class Transformer(nn.Module):
    def __init__(self,src_pad_idx,trg_pad_idx,enc_voc_size,dec_voc_size,max_len,d_model,n_head,d_hidden,n_layers,dropout_p,device):
        super(Transformer,self).__init__()

        self.encoder = Encoder(enc_voc_size,d_model,n_head,d_hidden,device,max_len,dropout_p,n_layers)
        self.decoder = Decoder(dec_voc_size,d_model,n_head,d_hidden,device,max_len,dropout_p,n_layers)

        self.src_pad_idx = src_pad_idx
        self.trg_pad_idx = trg_pad_idx
        self.device = device

    def make_casual_mask(self,q,k): 
        '''
        构建下三角掩码（time mask）
        '''
        len_q,len_k = q.size(1),k.size(1)
        mask = torch.tril(torch.ones((len_q,len_k),device=self.device)).type(torch.BoolTensor).to(self.device)
        return mask
    
    def make_pad_mask(self,q,k,pad_idx_q,pad_idx_k):
        '''
        构建padding掩码（src mask）
        SHAPE: (batch,time,len_q,len_k)
        '''
        len_q,len_k = q.size(1),k.size(1)
        q = q.ne(pad_idx_q).unsqueeze(1).unsqueeze(3)
        q = q.repeat(1,1,1,len_k) # (batch,1,len_q,len_k)
        k = k.ne(pad_idx_k).unsqueeze(1).unsqueeze(2)
        k = k.repeat(1,1,len_q,1) # (batch,1,len_q,len_k)
        mask = q & k
        return mask

    def forward(self,src,trg):
        src_mask = self.make_pad_mask(src,src,self.src_pad_idx,self.src_pad_idx)
        trg_mask = self.make_pad_mask(trg,trg,self.trg_pad_idx,self.trg_pad_idx) * self.make_casual_mask(trg,trg)
        src_trg_mask = self.make_pad_mask(trg,src,self.trg_pad_idx,self.src_pad_idx)


        enc = self.encoder(src,src_mask)
        out = self.decoder(trg,enc,trg_mask,src_trg_mask)
        return out

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    src_pad_idx = 1
    trg_pad_idx = 1
    enc_voc_size = 11
    dec_voc_size = 11
    max_len = 20
    d_model = 512
    n_head = 8
    d_hidden = 2048
    n_layers = 6
    dropout_p = 0.1

    model = Transformer(src_pad_idx,trg_pad_idx,enc_voc_size,dec_voc_size,max_len,d_model,n_head,d_hidden,n_layers,dropout_p,device).to(device)

    src = torch.tensor([[2,5,6,4,3,9,7,8,1,1,1],[3,4,5,6,7,8,9,1,1,1,1]],dtype=torch.long).to(device) # (batch,time)
    trg = torch.tensor([[2,7,4,3,5,9,8,1,1,1,1,1],[3,5,6,7,8,9,1,1,1,1,1,1]],dtype=torch.long).to(device) # (batch,time)

    out = model(src,trg[:,:-1])
    print(f"shape: {out.shape}\n输出:{out}")

Transformer#

多头注意力MHA#

Token和位置嵌入模块#

LayerNorm#

FFN前馈网络（别看名字，实际很简单）#

集成所有模块，构建完整的Transformer模型#