TR3：Pytorch复现Transformer

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍖 原作者：K同学啊

一、实验目的

• 从整体上把握Transformer模型，明白它是个什么东西，可以干嘛
• 读懂Transformer的复现代码

二、实验环境

• 语言环境：python 3.8
• 编译器：Jupyter notebook
• 深度学习环境：Pytorch
  • torch==2.4.0+cu124
  • torchvision==0.19.0+cu124

三、Transformer

Transformer 网络架构架构由 Ashish Vaswani 等人在 Attention Is All You Need一文中提出，并用于机器翻译任务，和以往网络架构有所区别的是，该网络架构中，编码器和解码器没有采用 RNN 或 CNN 等网络架构，而是采用完全依赖于注意力机制的架构。网络架构如下所示：


Transformer改进了RNN被人诟病的训练慢的特点，利用self-attention可以实现快速并行。

理解注意力机制

注意力机制是神经网络中一个迷人的概念，特别是在涉及到像 NLP 这样的任务时。它就像给模型一个聚光灯，让它能够集中注意力在输入序列的某些部分，同时忽略其他部分，就像我们人类在理解句子时关注特定的单词或短语一样。

现在，让我们深入了解一种特定类型的注意力机制，称为自注意力，也称为内部注意力。想象一下，当你阅读一句话时，你的大脑会自动突出显示重要的单词或短语来理解意思。这就是神经网络中自注意力的基本原理。它使序列中的每个单词都能“关注”其他单词，包括自己在内，以更好地理解上下文。

自注意力是如何工作的

嵌入：首先，模型将输入序列中的每个单词嵌入到一个高维向量表示中。这个嵌入过程允许模型捕捉单词之间的语义相似性。

查询、键和值向量：接下来，模型为序列中的每个单词计算三个向量：查询向量、键向量和值向量。在训练过程中，模型学习这些向量，每个向量都有不同的作用。查询向量表示单词的查询，即模型在序列中寻找的内容。键向量表示单词的键，即序列中其他单词应该注意的内容。值向量表示单词的值，即单词对输出所贡献的信息。

注意力分数：一旦模型计算了每个单词的查询、键和值向量，它就会为序列中的每一对单词计算注意力分数。这通常通过取查询向量和键向量的点积来实现，以评估单词之间的相似性。

SoftMax 归一化：然后，使用 softmax 函数对注意力分数进行归一化，以获得注意力权重。这些权重表示每个单词应该关注序列中其他单词的程度。注意力权重较高的单词被认为对正在执行的任务更为关键。

加权求和：最后，使用注意力权重计算值向量的加权和。这产生了每个序列中单词的自注意力机制输出，捕获了来自其他单词的上下文信息。

Transformer 模型的基础

在我们深入探讨Transformer模型的复杂工作原理之前，让我们花点时间欣赏其开创性的架构。正如我们之前讨论的，Transformer模型通过引入围绕自注意力机制的新颖方法，重塑了自然语言处理（NLP）的格局。

编码器-解码器架构：在Transformer的核心是其编码器-解码器架构——两个关键组件之间的共生关系，分别负责处理输入序列和生成输出序列。编码器和解码器中的每一层都包含相同的子层，包括自注意力机制和前馈网络。这种架构不仅有助于全面理解输入序列，而且能够生成上下文丰富的输出序列。

位置编码：尽管Transformer模型具有强大的功能，但它缺乏对元素顺序的内在理解——这是位置编码所解决的一个缺点。通过将输入嵌入与位置信息结合起来，位置编码使模型能够区分序列中元素的相对位置。这种细致的理解对于捕捉语言的时间动态和促进准确理解至关重要。

多头注意力：Transformer模型的一个显著特征是它能够同时关注输入序列的不同部分——这是多头注意力实现的。通过将查询、键和值向量分成多个头，并进行独立的自注意力计算，模型获得了对输入序列的细致透视，丰富了其表示，带有多样化的上下文信息。

前馈网络：与人类大脑能够并行处理信息的能力类似，Transformer模型中的每一层都包含一个前馈网络——一种能够捕捉序列中元素之间复杂关系的多功能组件。通过使用线性变换和非线性激活函数，前馈网络使模型能够在语言的复杂语义景观中航行，促进文本的稳健理解和生成。

更多 Transformer 知识可以参考文章：Transformer从零详细解读

四、复现Transformer

1. 多头注意力机制

import torch  
import torch.nn as nn  
import math
  
class MultiHeadAttention(nn.Module):  
    # n_heads：多头注意力的数量  
    # hid_dim：每个词输出的向量维度  
    def __init__(self, hid_dim, n_heads):  
        super(MultiHeadAttention, self).__init__()  
        self.hid_dim = hid_dim  
        self.n_heads = n_heads  
        # 强制 hid_dim 必须整除 h  
        assert hid_dim % n_heads == 0  
        # 定义 W_q 矩阵  
        self.w_q = nn.Linear(hid_dim, hid_dim)  
        # 定义 W_k 矩阵  
        self.w_k = nn.Linear(hid_dim, hid_dim)  
        # 定义 W_v 矩阵  
        self.w_v = nn.Linear(hid_dim, hid_dim)  
        self.fc  = nn.Linear(hid_dim, hid_dim)  
        # 缩放  
        self.scale = torch.sqrt(torch.FloatTensor([hid_dim // n_heads]))
        
    def forward(self, query, key, value, mask=None):  
        # 注意 Q，K，V的在句子长度这一个维度的数值可以一样，可以不一样。  
        # K: [64,10,300], 假设batch_size 为 64，有 10 个词，每个词的 Query 向量是 300 维
        # V: [64,10,300], 假设batch_size 为 64，有 10 个词，每个词的 Query 向量是 300 维
        # Q: [64,12,300], 假设batch_size 为 64，有 12 个词，每个词的 Query 向量是 300 维
        bsz = query.shape[0] 
        Q = self.w_q(query) 
        K = self.w_k(key) 
        V = self.w_v(value) 
        # 这里把 K Q V 矩阵拆分为多组注意力 
        # 最后一维就是是用 self.hid_dim // self.n_heads 来得到的，表示每组注意力的向量长度, 每个 head 的向量长度是：300/6=50
        # 64 表示 batch size，6 表示有 6组注意力，10 表示有 10 词，50 表示每组注意力的词的向量长度
        # K: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # V: [64,10,300] 拆分多组注意力 -> [64,10,6,50] 转置得到 -> [64,6,10,50]
        # Q: [64,12,300] 拆分多组注意力 -> [64,12,6,50] 转置得到 -> [64,6,12,50]
        # 转置是为了把注意力的数量 6 放到前面，把 10 和 50 放到后面，方便下面计算   
        Q = Q.view(bsz, -1, self.n_heads, self.hid_dim //   
                   self.n_heads).permute(0, 2, 1, 3)   
        K = K.view(bsz, -1, self.n_heads, self.hid_dim //   
                   self.n_heads).permute(0, 2, 1, 3)   
        V = V.view(bsz, -1, self.n_heads, self.hid_dim //   
                   self.n_heads).permute(0, 2, 1, 3)
                      
        # 第 1 步：Q 乘以 K的转置，除以scale   
        # [64,6,12,50] * [64,6,50,10] = [64,6,12,10]   
        # attention：[64,6,12,10]   
        attention = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale
        
        # 如果 mask 不为空，那么就把 mask 为 0 的位置的 attention 分数设置为 -1e10，这里用“0”来指示哪些位置的词向量不能被attention到，比如padding位置，当然也可以用“1”或者其他数字来指示，主要设计下面2行代码的改动。
        if mask is not None:   
            attention = attention.masked_fill(mask == 0, -1e10) 
              
            # 第 2 步：计算上一步结果的 softmax，再经过 dropout，得到 attention。
            # 注意，这里是对最后一维做 softmax，也就是在输入序列的维度做 softmax   
            # attention: [64,6,12,10]   
        attention = torch.softmax(attention, dim=-1) 
          
        # 第3步，attention结果与V相乘，得到多头注意力的结果   
        # [64,6,12,10] * [64,6,10,50] = [64,6,12,50]   
        # x: [64,6,12,50]   
        x = torch.matmul(attention, V) 
          
        # 因为 query 有 12 个词，所以把 12 放到前面，把 50 和 6 放到后面，方便下面拼接多组的结果
        # x: [64,6,12,50] 转置-> [64,12,6,50]   
        x = x.permute(0, 2, 1, 3).contiguous()   
        # 这里的矩阵转换就是：把多组注意力的结果拼接起来   
        # 最终结果就是 [64,12,300]   
        # x: [64,12,6,50] -> [64,12,300]   
        x = x.view(bsz, -1, self.n_heads * (self.hid_dim // self.n_heads))
        x = self.fc(x) 
        return x 

2. 前馈传播

class Feedforward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(Feedforward, self).__init__()
        # 两层线性映射和激活函数ReLU
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
    def forward(self, x):
        x = torch.nn.functional.relu(self.linear1(x))
        x = self.dropout(x)
        x = self.linear2(x)
        return x

3. 位置编码

class PositionalEncoding(nn.Module):
    "实现位置编码"
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        # 初始化Shape为(max_len, d_model)的PE (positional encoding)
        pe = torch.zeros(max_len, d_model).to(device)
        
        # 初始化一个tensor [[0, 1, 2, 3, ...]]
        position = torch.arange(0, max_len).unsqueeze(1)
        # 这里就是sin和cos括号中的内容，通过e和ln进行了变换
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term) # 计算PE(pos, 2i)
        pe[:, 1::2] = torch.cos(position * div_term) # 计算PE(pos, 2i+1)
        
        pe = pe.unsqueeze(0) # 为了方便计算，在最外面在unsqueeze出一个batch
        
        # 如果一个参数不参与梯度下降，但又希望保存model的时候将其保存下来
        # 这个时候就可以用register_buffer
        self.register_buffer("pe", pe)
    def forward(self, x):
        """
        x 为embedding后的inputs，例如(1,7, 128)，batch size为1,7个单词，单词维度为128
        """
        # 将x和positional encoding相加。
        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
        return self.dropout(x)

4. 解码层

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        # 编码器层包含自注意力机制和前馈神经网络
        self.self_attn   = MultiHeadAttention(d_model, n_heads)
        self.feedforward = Feedforward(d_model, d_ff, dropout)
        self.norm1   = nn.LayerNorm(d_model)
        self.norm2   = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x, mask):
        # 自注意力机制
        attn_output = self.self_attn(x, x, x, mask)
        x = x + self.dropout(attn_output)
        x = self.norm1(x)
        # 前馈神经网络
        ff_output = self.feedforward(x)
        x = x + self.dropout(ff_output)
        x = self.norm2(x)
        return x

5. 编码层

class DecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        # 解码器层包含自注意力机制、编码器-解码器注意力机制和前馈神经网络
        self.self_attn   = MultiHeadAttention(d_model, n_heads)
        self.enc_attn    = MultiHeadAttention(d_model, n_heads)
        self.feedforward = Feedforward(d_model, d_ff, dropout)
        self.norm1   = nn.LayerNorm(d_model)
        self.norm2   = nn.LayerNorm(d_model)
        self.norm3   = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x, enc_output, self_mask, context_mask):
        # 自注意力机制
        attn_output = self.self_attn(x, x, x, self_mask)
        x           = x + self.dropout(attn_output)
        x           = self.norm1(x)
        # 编码器-解码器注意力机制
        attn_output = self.enc_attn(x, enc_output, enc_output, context_mask)
        x           = x + self.dropout(attn_output)
        x           = self.norm2(x)
        # 前馈神经网络
        ff_output = self.feedforward(x)
        x = x + self.dropout(ff_output)
        x = self.norm3(x)
        return x

6. Transformer模型构建

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, n_encoder_layers, n_decoder_layers, d_ff, dropout=0.1):
        super(Transformer, self).__init__()
        # Transformer 模型包含词嵌入、位置编码、编码器和解码器
        self.embedding           = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
        self.encoder_layers      = nn.ModuleList([EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_encoder_layers)])
        self.decoder_layers      = nn.ModuleList([DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_decoder_layers)])
        self.fc_out              = nn.Linear(d_model, vocab_size)
        self.dropout             = nn.Dropout(dropout)
    def forward(self, src, trg, src_mask, trg_mask):
        # 词嵌入和位置编码
        src = self.embedding(src)
        src = self.positional_encoding(src)
        trg = self.embedding(trg)
        trg = self.positional_encoding(trg)
        # 编码器
        for layer in self.encoder_layers:
            src = layer(src, src_mask)
        # 解码器
        for layer in self.decoder_layers:
            trg = layer(trg, src, trg_mask, src_mask)
        # 输出层
        output = self.fc_out(trg)
        return output

7. 使用示例

vocab_size = 10000  # 假设词汇表大小为10000
d_model    = 128
n_heads    = 8
n_encoder_layers = 6
n_decoder_layers = 6
d_ff             = 2048
dropout          = 0.1
device = torch.device('cpu')
 
transformer_model = Transformer(vocab_size, d_model, n_heads, n_encoder_layers, n_decoder_layers, d_ff, dropout)
# 定义输入，这里的输入是假设的，需要根据实际情况修改
src = torch.randint(0, vocab_size, (32, 10))  # 源语言句子
trg = torch.randint(0, vocab_size, (32, 20))  # 目标语言句子
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)  # 掩码，用于屏蔽填充的位置
trg_mask = (trg != 0).unsqueeze(1).unsqueeze(2)  # 掩码，用于屏蔽填充的位置
# 模型前向传播
output = transformer_model(src, trg, src_mask, trg_mask)
print(output.shape)

五、总结

通过本次 Transformer 的学习与模型复现，我对其架构和原理有了初步的了解。

在理论层面，明白了 Transformer 基于注意力机制构建，编码器 - 解码器架构、位置编码、多头注意力和前馈网络是其核心要素。注意力机制让模型能聚焦关键信息，多头注意力从多维度捕捉上下文，位置编码赋予位置感知，前馈网络处理复杂关系。

在实践方面，实现了 PyTorch 的复现过程。从多头注意力机制、前馈传播、位置编码、编码层和解码层的实现，到最终构建Transformer模型并给出使用示例。这有助于理解模型各部分如何协同工作，以及如何将理论转化为实际代码，为进一步应用和优化Transformer模型奠定了基础，也为深入研究深度学习中的序列处理模型提供了重要参考。

Transformer模型作为一种强大的深度学习架构，在众多自然语言处理和其他序列数据相关的实际问题中都有广泛应用，我主要关注其在医学领域中结合自然语言处理和计算机视觉的常见应用：

自然语言处理在医学领域的应用

• 医学文献检索与信息提取：利用Transformer的语言理解能力，可以更准确地对医学文献进行检索和信息提取。例如，研究人员可以输入特定的医学问题或关键词，Transformer模型能够理解问题的语义，从大量的医学文献中筛选出与之相关的高质量文献，并提取其中的关键信息，如疾病的病因、症状、治疗方法等，帮助医学工作者快速获取所需知识，节省时间和精力。
• 电子病历分析与处理：Transformer模型可用于分析电子病历中的文本数据，包括患者的症状描述、检查结果、诊断结论、治疗过程等。通过对这些文本的理解和分析，模型可以帮助医生更好地了解患者的病情历史，发现潜在的疾病模式和风险因素，为诊断和治疗提供参考依据。此外，还可以实现电子病历的自动编码和分类，提高医疗数据管理的效率和准确性。
• 医学问答系统：基于Transformer构建的医学问答系统能够理解用户提出的自然语言问题，并生成准确、有用的答案。无论是患者咨询常见的健康问题，还是医学专业人员探讨复杂的病例，问答系统都可以快速提供相关的医学知识和建议，提高医疗信息的可及性和交流效率.
• 医学文本生成：可以生成医学报告、疾病诊断书、治疗方案等文本内容。例如，根据患者的检查数据和基本信息，Transformer模型能够生成初步的诊断报告，为医生提供参考，辅助其进行更准确的诊断和治疗决策。同时，在医学教育领域，也可以用于生成教学案例、知识点讲解等文本材料，帮助学生更好地学习和理解医学知识.

计算机视觉在医学领域的应用

• 医学图像分割：医学图像分割是指将医学图像中的不同组织、器官或病变区域等进行精确划分，这对于疾病的诊断、治疗计划制定以及手术导航等都具有重要意义。Transformer的全局建模能力使其能够有效地捕捉图像中的长距离依赖关系，从而更准确地分割出医学图像中的目标区域，例如将CT图像中的肿瘤、MRI图像中的脑部组织等进行精确分割，为医生提供更详细、准确的图像信息.
• 医学图像分类：可用于对医学图像进行分类，如判断X光片中是否存在骨折、CT图像中是否有肺部疾病等。通过学习大量的医学图像数据，Transformer模型能够自动提取图像的特征，并根据这些特征对图像进行准确分类，辅助医生进行快速的疾病筛查和诊断.
• 医学图像重建与增强：在医学成像过程中，由于设备限制、患者运动等因素，可能会导致图像质量下降。Transformer模型可以学习图像的潜在特征表示，对低质量的医学图像进行重建和增强，提高图像的清晰度和分辨率，有助于医生更准确地观察和分析图像中的细节信息，如从欠采样的MRI图像中恢复高分辨率图像.
• 医学图像配准：医学图像配准是将不同时间、不同模态或不同视角下获取的医学图像进行对齐的过程，对于多模态图像融合、疾病监测和治疗评估等具有重要价值。Transformer能够捕获图像中的全局上下文信息和长程依赖关系，更好地处理大的空间变形和多模态特征融合问题，从而实现更准确的医学图像配准.
• 医学图像合成：可以根据给定的条件或先验知识生成逼真的医学图像，如生成特定疾病的模拟图像、合成缺失模态的医学图像等。这有助于医生更好地理解疾病的形态和特征，进行术前规划和模拟手术，同时也为医学研究提供了更多的数据资源.