【AI小项目4】用Pytorch从头实现Transformer（详细注解）

AI完全体

已于 2024-10-15 23:40:53 修改

阅读量1.8k

点赞数 21

分类专栏： AI小项目文章标签：人工智能 pytorch transformer 深度学习机器学习注意力机制语言模型

于 2024-09-13 23:12:01 首次发布

本文链接：https://blog.csdn.net/weixin_43221845/article/details/142104186

版权

【AI论文解读】【AI知识点】【AI小项目】【AI战略思考】

一、前期准备工作

学习如何读AI论文

看了B站李沐老师的下面几个视频

如何读论文【论文精读·1】
9年后重读深度学习奠基作之一：AlexNet【论文精读·2】
AlexNet论文逐段精读【论文精读】

读Transformer原始论文

我是看B站李沐老师的Transformer论文精读视频，结合原作论文和chatgpt来学习的。chatgpt用于查一些背景知识。

原作论文作为参考，自己读当然OK，但是如果有大神帮你解读，对于小白来说，站在巨人的肩膀上效率可能更高，比如大神会提到经典论文中哪些部分是现在还常用的，哪些部分已经过时或者不重要。

李沐老师的视频：Transformer论文逐段精读【论文精读】
原作论文：Attention Is All You Need

用Pytorch从头实现Transformer

去kaggle或者github上找一份不错的范例代码来参考学习

我参考的代码kaggle地址：Transformer from scratch using pytorch

二、我的完整代码实现

1.kaggle上有我的完整代码，这里贴一份是为了方便那些不会科学上网打不开kaggle的读者。
2.代码有不懂的地方，可以看下背景知识补充部分

我的kaggle代码地址

先看一下论文中完整的transformer架构图：

1.导入库

import torch.nn as nn
import torch
import torch.nn.functional as F
import math,copy,re
import warnings
import pandas as pd
import numpy as np
import seaborn as sns
import torchtext
import matplotlib.pyplot as plt
warnings.simplefilter("ignore")
print(torch.__version__)

2.基本组件

创建词嵌入

这个 Embedding 类用于将词汇索引转换为相应的嵌入向量，以便输入到神经网络中。
nn.Embedding 层通过查表的方式，将每个词汇的索引映射到一个固定维度的嵌入向量上。

class Embedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        """
        Args:
            vocab_size: 词汇表大小
            embed_dim: 嵌入维度
        """
        super(Embedding, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
    def forward(self, x):
        """
        Args:
            x: 输入向量，通常是一个包含词汇索引的张量（整数表示词汇在词汇表中的位置）
            out: 嵌入向量
        """
        out = self.embed(x)
        return out

位置嵌入

这个 PositionalEmbedding 模块通过对每个位置进行正弦和余弦编码，为每个输入序列的嵌入向量提供了位置信息，使模型能够感知序列中词汇的顺序。

class PositionalEmbedding(nn.Module):
    def __init__(self,max_seq_len,embed_model_dim):
        """
        Args:
            max_seq_len: 输入序列的最大长度
            embed_model_dim: 嵌入的维度
        """
        super(PositionalEmbedding, self).__init__()
        self.embed_dim = embed_model_dim

        # 生成位置嵌入矩阵
        # 为每个位置计算对应的正弦和余弦编码
        pe = torch.zeros(max_seq_len,self.embed_dim)
        for pos in range(max_seq_len):
            for i in range(0,self.embed_dim,2):
                # i = 0,2,4,6 ...
                pe[pos, i] = np.sin(pos/(10000**(i/self.embed_dim)))
                if i+1 < self.embed_dim:
                    pe[pos, i+1] = np.cos(pos/(10000**(i/self.embed_dim)))
                    
        # 扩展维度以适应批量输入
        # 使用unsqueeze(0)在第一个维度上增加一个维度，使pe的形状变为(1, max_seq_len, embed_dim)。
        pe = pe.unsqueeze(0)
        
        # 将 pe 注册为 buffer，防止它在训练过程中被更新
        # 缓冲区的作用：在模型保存和加载时，缓冲区会一同保存，但不会作为模型参数参与训练。
        self.register_buffer('pe', pe)


    def forward(self, x):
        """
        Args:
            x: 输入的词嵌入张量，形状通常为(batch_size, seq_len, embed_dim)
        Returns:
            x: 输出
        """
      
        # 使输入嵌入向量的值相对较大
        x = x * math.sqrt(self.embed_dim)
        # 获取输入序列的长度
        seq_len = x.size(1)
        # 提取对应长度的位置嵌入，将位置嵌入与词嵌入相加，融合位置信息
        x = x + torch.autograd.Variable(self.pe[:, :seq_len, :], requires_grad=False)
        return x

自注意力

这个 MultiHeadAttention 类实现了多头注意力机制（Multi-Head Attention），是 Transformer 中的一个重要组件。多头注意力机制允许模型通过多个注意力头（heads）从不同的子空间中关注不同的部分。

论文中多头注意力机制的图

在这里插入图片描述

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim=512, n_heads=8):
        """
        Args:
            embed_dim: 嵌入向量的维度
            n_heads: 注意力头的数量
        """
        super(MultiHeadAttention, self).__init__()

        self.embed_dim = embed_dim  # 512维嵌入向量
        self.n_heads = n_heads  # 8个注意力头
        self.single_head_dim = int(self.embed_dim / self.n_heads)  # 每个注意力头的维度 = 512 / 8 = 64

        # 定义线性变换矩阵，用于生成 Query, Key 和 Value 矩阵
        self.query_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)
        self.key_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)
        self.value_matrix = nn.Linear(self.single_head_dim, self.single_head_dim, bias=False)

        # 输出线性层，作用是将多头注意力的输出拼接并变换回原始嵌入维度
        self.out = nn.Linear(self.n_heads * self.single_head_dim, self.embed_dim)


    def forward(self, key, query, value, mask=None):
        """
        Args:
           key: Key 向量
           query: Query 向量
           value: Value 向量
           mask: 用于屏蔽不需要计算注意力的部分
        Returns:
           output: 多头注意力机制的输出
        """
        batch_size = key.size(0)
        seq_length = key.size(1)
        seq_length_query = query.size(1)

        # 重塑 key, query 和 value 的维度，使其适应多头注意力的格式
        # key、query、value 的输入维度为 (batch_size, seq_length, embed_dim)，例如 (32, 10, 512) 表示 32 个序列，每个序列长度为 10，每个词的嵌入维度为 512。
        # 使用view()而不是reshape()的理由：view() 不会创建新张量，只是修改现有张量的元数据（假设内存布局是连续的），因此速度快，内存开销低。
        key = key.view(batch_size, seq_length, self.n_heads, self.single_head_dim)  # (32, 10, 8, 64)
        query = query.view(batch_size, seq_length_query, self.n_heads, self.single_head_dim)
        value = value.view(batch_size, seq_length, self.n_heads, self.single_head_dim)

        # 线性变换得到 Q, K, V 矩阵
        k = self.key_matrix(key)  # (32, 10, 8, 64)
        q = self.query_matrix(query)  # (32, 10, 8, 64)
        v = self.value_matrix(value)  # (32, 10, 8, 64)

        # 交换维度，适应矩阵乘法的形状
        q = q.transpose(1, 2)  # (32, 8, 10, 64)
        k = k.transpose(1, 2)  # (32, 8, 10, 64)
        v = v.transpose(1, 2)  # (32, 8, 10, 64)

        # 计算注意力分数
        k_adjusted = k.transpose(-1, -2)  # 调整 Key 矩阵以便于与 Query 相乘 (32, 8, 64, 10)
        product = torch.matmul(q, k_adjusted)  # 矩阵乘法得到注意力分数 (32, 8, 10, 10)

        # 如果提供了 mask，则屏蔽某些位置的注意力分数
        if mask is not None:
            product = product.masked_fill(mask == 0, float("-1e20"))

        # 对 Key 的维度进行缩放，缓解梯度消失问题
        product = product / math.sqrt(self.single_head_dim)  # 缩放因子 √64

        # 计算注意力权重
        scores = F.softmax(product, dim=-1)  # (32, 8, 10, 10)

        # 将注意力权重与 Value 相乘，得到加权后的输出
        scores = torch.matmul(scores, v)  # (32, 8, 10, 64)

        # 将多头的输出拼接起来
        concat = scores.transpose(1, 2).contiguous().view(batch_size, seq_length_query, self.single_head_dim * self.n_heads)  # (32, 10, 512)

        # 通过线性层输出
        output = self.out(concat)  # (32, 10, 512)

        return output

3.编码器

论文中编码器的图

TransformerBlock 是 Transformer 编码器的基本单元，它由多头注意力机制、残差连接、归一化层、前馈神经网络和 dropout 组成。

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, expansion_factor=4, n_heads=8):
        super(TransformerBlock, self).__init__()
        
        """
        Args:
           embed_dim: 嵌入向量的维度
           expansion_factor: 前馈网络的扩展因子，用于扩展中间层的维度
           n_heads: 注意力头的数量
        """
        self.attention = MultiHeadAttention(embed_dim, n_heads)  # 多头注意力机制
        
        self.norm1 = nn.LayerNorm(embed_dim)  # 第一个归一化层（LayerNorm）
        self.norm2 = nn.LayerNorm(embed_dim)  # 第二个归一化层
        
        self.feed_forward = nn.Sequential(  # 前馈神经网络部分
                          nn.Linear(embed_dim, expansion_factor * embed_dim),  # 线性层，维度从 embed_dim 扩展到 expansion_factor * embed_dim
                          nn.ReLU(),  # 激活函数，引入非线性
                          nn.Linear(expansion_factor * embed_dim, embed_dim)  # 再次将维度缩小到原始的 embed_dim
        )

        self.dropout1 = nn.Dropout(0.2)  # 第一个 dropout，用于防止过拟合
        self.dropout2 = nn.Dropout(0.2)  # 第二个 dropout

    def forward(self, key, query, value):
        """
        Args:
           key: Key 向量
           query: Query 向量
           value: Value 向量
        
        Returns:
           norm2_out: 经过 Transformer Block 的输出
        """
        
        attention_out = self.attention(key, query, value)  # 计算多头注意力输出
        attention_residual_out = attention_out + value  # 残差连接，添加 Value 向量
        norm1_out = self.dropout1(self.norm1(attention_residual_out))  # 归一化 + dropout

        feed_fwd_out = self.feed_forward(norm1_out)  # 前馈网络计算
        feed_fwd_residual_out = feed_fwd_out + norm1_out  # 残差连接
        norm2_out = self.dropout2(self.norm2(feed_fwd_residual_out))  # 归一化 + dropout

        return norm2_out

TransformerEncoder 类由多个 TransformerBlock 组成，它负责对输入序列进行编码。输入经过词嵌入层和位置编码层后，进入多个 TransformerBlock 层进行处理。

class TransformerEncoder(nn.Module):
    """
    Transformer 编码器
    
    Args:
        seq_len : 输入序列的长度
        vocab_size: 词汇表大小
        embed_dim: 嵌入向量的维度
        num_layers: 编码器的层数
        expansion_factor: 前馈网络中的扩展因子
        n_heads: 多头注意力机制中的头数
        
    Returns:
        out: 编码器的输出
    """
    def __init__(self, seq_len, vocab_size, embed_dim, num_layers=2, expansion_factor=4, n_heads=8):
        super(TransformerEncoder, self).__init__()
        
        self.embedding_layer = Embedding(vocab_size, embed_dim)  # 词嵌入层
        self.positional_encoder = PositionalEmbedding(seq_len, embed_dim)  # 位置编码

        # 多层 Transformer Block
        self.layers = nn.ModuleList([TransformerBlock(embed_dim, expansion_factor, n_heads) for i in range(num_layers)])
    
    def forward(self, x):
        embed_out = self.embedding_layer(x)  # 词嵌入
        out = self.positional_encoder(embed_out)  # 加上位置编码
        for layer in self.layers:
            out = layer(out, out, out)  # 通过多个 TransformerBlock

        return out  # 输出维度为 (batch_size, seq_len, embed_dim)

4.解码器

论文中解码器的图

这个 DecoderBlock 类是 Transformer 解码器（Decoder）中的一个基本模块，主要用于处理输入的解码序列，同时结合编码器（Encoder）输出的信息。它包含了多头自注意力机制、前馈神经网络和残差连接等部分。

class DecoderBlock(nn.Module):
    def __init__(self, embed_dim, expansion_factor=4, n_heads=8):
        super(DecoderBlock, self).__init__()

        """
        Args:
           embed_dim: 嵌入向量的维度
           expansion_factor: 前馈网络中的扩展因子
           n_heads: 多头注意力头的数量
        """
        self.attention = MultiHeadAttention(embed_dim, n_heads=n_heads)  # 解码器中的自注意力机制
        self.norm = nn.LayerNorm(embed_dim)  # 归一化层，用于稳定网络训练
        self.dropout = nn.Dropout(0.2)  # dropout，用于防止过拟合
        self.transformer_block = TransformerBlock(embed_dim, expansion_factor, n_heads)  # 解码器中的 transformer 模块

    def forward(self, key, query, x, mask):
        """
        Args:
           key: Key 向量，通常来自编码器的输出
           query: Query 向量，来自解码序列
           x: 输入解码序列
           mask: 遮掩掩码，用于屏蔽未来信息

        Returns:
           out: 解码器块的输出
        """

        # 对解码器输入进行自注意力计算，传入 mask 防止查看未来信息
        attention = self.attention(x, x, x, mask=mask)  # 输出大小为 (batch_size, seq_len, embed_dim)

        # 残差连接，并通过归一化和 dropout
        value = self.dropout(self.norm(attention + x))

        # 将处理后的 `value` 传入 TransformerBlock，与来自编码器的 `key` 和 `query` 进行进一步处理
        out = self.transformer_block(key, query, value)

        return out

这个 TransformerDecoder 类是 Transformer 解码器的完整实现，它由词嵌入层、位置编码层、多个解码块（DecoderBlock），以及最终的全连接层组成。解码器的主要作用是生成输出序列，例如机器翻译中的目标语言序列。

class TransformerDecoder(nn.Module):
    # 初始化方法定义了解码器的主要组件，包括嵌入层、位置编码层、多个 DecoderBlock 层和一个全连接层。
    def __init__(self, target_vocab_size, embed_dim, seq_len, num_layers=2, expansion_factor=4, n_heads=8):
        super(TransformerDecoder, self).__init__()
        """  
        Args:
           target_vocab_size: 目标词汇表的大小
           embed_dim: 嵌入向量的维度
           seq_len : 输入序列的长度
           num_layers: 解码器层的数量
           expansion_factor: 前馈网络中的扩展因子
           n_heads: 多头注意力头的数量
        """
        
        # 词嵌入层
        self.word_embedding = nn.Embedding(target_vocab_size, embed_dim)
        
        # 位置编码层
        self.position_embedding = PositionalEmbedding(seq_len, embed_dim)

        # 多个解码层（DecoderBlock）
        self.layers = nn.ModuleList(
            [
                DecoderBlock(embed_dim, expansion_factor=expansion_factor, n_heads=n_heads) 
                for _ in range(num_layers)
            ]
        )
        
        # 输出层，全连接层将嵌入维度转换为词汇表大小
        self.fc_out = nn.Linear(embed_dim, target_vocab_size)
        
        # dropout 用于防止过拟合
        self.dropout = nn.Dropout(0.2)

    # 前向传播过程中，输入序列首先经过嵌入层和位置编码层，然后通过多个 DecoderBlock 进行处理，最后通过全连接层输出每个位置上词汇的概率分布。
    def forward(self