Transformer 模型 pytorch 复现代码+注解+推理测试

https://orzzz.net/directory/leetcode/transformer/index.html

'''
# System --> Windows & Python3.10.0
# File ----> main.py
# Author --> Illusionna
# Create --> 2024/07/10 19:32:15
'''
# -*- Encoding: UTF-8 -*-


"""
Step 1: 查看当前 GPU 的所有基础信息
    >>> nvidia-smi
    CUDA Version: 12.1

Step 2: 配置 main.py 程序依赖环境
    >>> conda create -n NLP python==3.10.0
    >>> conda activate NLP
    >>> pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

Step 3: 激活并测试环境
    >>> conda activate NLP
    >>> python
    >>> import torch
    >>> print(torch.__version__)
    2.3.1+cu121
    >>> print(torch.cuda.is_available())
    True
    >>> exit(0)

Step 4: Transformer 简单测试示例推理.
    >>> python main.py
"""

"""
"Attention is All You Need", 代码 pytorch 推荐两篇:
1. 哈佛大学 NLP 研究组
    旧版: http://nlp.seas.harvard.edu/2018/04/03/attention.html
    新版: https://nlp.seas.harvard.edu/annotated-transformer
2. 台湾小哥通俗代码
    https://github.com/jadore801120/attention-is-all-you-need-pytorch
"""


import math
import copy
import torch


def debug(**kwargs) -> None:
    """
    调试打印, 仅此处局部 import 库.

    >>> debug(kwargs='2.718281828')
    """
    # <S>---------------------------------------------------------------
    import os
    import sys
    import random
    import inspect
    # ---------------------------------------------------------------<E>
    os.system('')
    line = inspect.getframeinfo(sys._getframe(1))
    file = os.path.relpath(line.filename, os.getcwd())
    print('<S>---------------------------------------------------------------')
    print(f'\033[3{random.randint(1, 6)}m[+debug] "{file}", line {line.lineno}')
    for key, value in kwargs.items():
        print('')
        print(key, '=')
        print(value, end='\n')
    print('\033[0m', end='')
    print('---------------------------------------------------------------<E>')


def clone(module: object, N: int) -> object:
    """
    克隆 N 个相同的块层.
    """
    return torch.nn.ModuleList([copy.deepcopy(module) for _ in range(0, N, 1)])


def subsequentMask(size: int) -> torch.Tensor:
    """
    掩码矩阵图示见 https://s21.ax1x.com/2024/07/10/pkf52UH.png

    创建掩码矩阵, 屏蔽后续位置, 防止解码器计算自注意力"看到"未来的词, 保持序列生成的因果关系, 自注意力是并行计算, 掩码设计可以让解码器训练阶段高效处理长序列数据, 同时保持正确序列生成顺序.
    >>> subsequentMask(4)
    >>> tensor(
        [
            [
                [ True, False, False, False],
                [ True, True, False, False],
                [ True, True, True, False],
                [ True, True, True, True]
            ]
        ]
    )
    """
    attention_shape = (1, size, size)
    # torch.triu 把一个矩阵强制转化成上三角矩阵, 从索引为 1 的对角线处处理.
    subsequent_mask = torch.triu(torch.ones(attention_shape), diagonal=1).type(torch.uint8)
    return subsequent_mask == 0


def scaledDotProductAttention(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    attention_mask: torch.Tensor,
    dropout: object
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    计算注意力, 返回得分和权重, 函数公式见 https://s21.ax1x.com/2024/07/10/pkfThG9.png
    >>> scores = [
        [0.1, 0.2, 0.3, 0.4, 0.5],
        [0.6, 0.7, 0.8, 0.9, 1.0],
        [1.1, 1.2, 1.3, 1.4, 1.5],
        [1.6, 1.7, 1.8, 1.9, 2.0],
        [2.1, 2.2, 2.3, 2.4, 2.5]
    ]
    # mask 矩阵也可以为布尔矩阵, 效果是等价于 0-1 矩阵.
    >>> mask = [
        [1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0]
    ]
    >>> scores.masked_fill --> [
        [0.1, 0.2, 0.3, -INFINITY, -INFINITY],
        [0.6, 0.7, 0.8, -INFINITY, -INFINITY],
        [1.1, 1.2, 1.3, -INFINITY, -INFINITY],
        [1.6, 1.7, 1.8, -INFINITY, -INFINITY],
        [2.1, 2.2, 2.3, -INFINITY, -INFINITY]
    ]
    """
    INFINITY = 1e9
    d_k = Q.size(dim=-1)
    # 除以根号 d_k 即 scale 缩放, 原因在于防止 softmax 推向梯度极小的区域不利于训练.
    scores = torch.matmul(Q, K.transpose(dim0=-1, dim1=-2)) / math.sqrt(d_k)
    if attention_mask is not None:
        # 遮掩处值全部置为负无穷, 以便下一行 softmax 函数计算其为数值零.
        scores = scores.masked_fill(attention_mask == False, -INFINITY)
    # 注意力权重概率表示, Q 矩阵元素对应 K 矩阵元素的相似性或重要性, 正相关.
    probability_attention = scores.softmax(dim=-1)
    # 丢弃一部分神经元.
    if dropout is not None:
        probability_attention = dropout(probability_attention)
    # 除了 Scaled Dot Product Attention 点积注意力以外, 还有一种 Additive Attention 加性注意力, 但斯坦福自然语言处理研究组发现, 在实践中, 点积注意力更快, 更节省空间, 因为它可以使用底层库的高度优化的矩阵乘法 torch.matmul() 实现.
    return (torch.matmul(probability_attention, V), probability_attention)


class MultiHeadAttention(torch.nn.Module):
    """
    多头(交叉)注意力, 架构见 https://s21.ax1x.com/2024/07/10/pkhVUXt.png
    """
    def __init__(self, *args, h: int, d_model: int, dropout: float=0.1, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 若等号成立, 则继续执行, 否则抛出 AssertionError 错误.
        assert d_model % h == 0
        self.d_k = d_model // h
        self.n_heads = h
        # 克隆 4 块.
        self.linears = clone(torch.nn.Linear(d_model, d_model), 4)
        self.dropout = torch.nn.Dropout(dropout)

    def forward(
        self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
        mask: torch.Tensor | None = None
    ) -> torch.Tensor:
        if mask is not None:
            # 所有的 h 个头都应用掩码.
            mask = mask.unsqueeze(dim=1)
        # RESIDUAL = Q
        batch_size = Q.size(dim=0)
        # <S>---------------------------------------------------------------
        # 循环三次, linear 每次取自 self.linears 迭代器, x 张量每次取自列表 [Q, K, V] 中的元素.
        [Q, K, V] = [
            # view() 相对于重塑, batch_size 批量, -1 代表自动根据 h * d_k 进行调整.
            linear(x).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
            for linear, x in zip(self.linears, [Q, K, V])
        ]
        '''
        # 上面列表推导式等价于下面写法.
        Q = self.linears[0](Q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = self.linears[1](K).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.linears[2](V).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        '''
        # ---------------------------------------------------------------<E>
        (x, _) = scaledDotProductAttention(
            Q = Q, K = K, V = V,
            attention_mask = mask, dropout = self.dropout
        )
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_k)
        del Q; del K; del V
        # <S>---------------------------------------------------------------
        # 最后一个块进行推理, 即索引为 3 的块.
        '''
        # 台湾小哥在最后返回处进行了残存网络连接, 然后再正则化.
        norm = LayerNorm(x.size(dim=-1))
        return norm(self.linears[-1](x) + RESIDUAL)
        '''
        return self.linears[-1](x)
        # ---------------------------------------------------------------<E>


class PositionWiseFeddForward(torch.nn.Module):
    """
    位置逐元素全连接前馈神经网络, 函数公式见 https://s21.ax1x.com/2024/07/10/pkfvZr9.png
    """
    def __init__(
        self,
        *args,
        d_model: int,
        d_ff: int,
        dropout: float = 0.1,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.W1 = torch.nn.Linear(d_model, d_ff, bias=True)
        self.W2 = torch.nn.Linear(d_ff, d_model, bias=True)
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # torch.Tensor 含有属性 relu() 函数, 可以直接调用.
        # return self.W2(self.dropout(self.W1(x).relu()))
        return self.W2(self.dropout(torch.relu(self.W1(x))))


class PositionalEncoding(torch.nn.Module):
    """
    位置编码, 原理见该作者博客 https://wmathor.com/index.php/archives/1453
    """
    def __init__(
        self,
        *args,
        d_model: int,
        dropout: float = 0.1,
        max_tokens: int = 5000,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.dropout = torch.nn.Dropout(dropout)
        pe = torch.zeros(max_tokens, d_model)
        # 创建 max_tokens x 1 的列向量.
        position = torch.arange(0, max_tokens, 1).unsqueeze(dim=1)
        '''
        矩阵的所有行, 每行从索引 0 开始, 跳步 2 单位.
        >>> matrix = torch.Tensor(
            [
                [1, 2, 3, 4, 5, 6, 7, 8, 9],
                [2, 2, 3, 4, 5, 6, 7, 8, 9],
                [3, 2, 3, 4, 5, 6, 7, 8, 9]
            ]
        )
        >>> matrix[:, 0::2] = torch.Tensor(
            [1, 3, 5, 7, 9],
            [2, 3, 5, 7, 9],
            [3, 3, 5, 7, 9]
        )
        >>> matrix[:, 0:-1:2] = torch.Tensor(
            [1, 3, 5, 7],
            [2, 3, 5, 7],
            [3, 3, 5, 7]
        )
        '''
        divisor = torch.exp(torch.arange(0, d_model, 2) * -(math.log(1e4) / d_model))
        pe[:, 0::2] = torch.sin(position * divisor)
        pe[:, 1::2] = torch.cos(position * divisor)
        '''
        >>> pe = torch.Tensor(
            [
                [1, 2, 3],
                [4, 5, 6],
                [7, 8, 9]
            ]
        )
        >>> pe.unsqueeze(0) = torch.Tensor(
            [
                [
                    [1, 2, 3],
                    [4, 5, 6],
                    [7, 8, 9]
                ]
            ]
        )
        # 增加维度, 给索引为 0 维即 batch 维增加一单位, 相当于添加一个大小为 1 的 batch 维.
        '''
        pe = pe.unsqueeze(dim=0)
        # 将 pe 注册成为模型的缓冲区, 可以在 forward 函数中访问.
        self.register_buffer('pe', pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        1. 缓冲区是模型的一部分;
        2. 位置编码是固定不变的, 训练过程不会更新的, 所以将 pe 注册到缓冲区是合适的;
        3. 梯度计算设置为 False, 意味反向传播计算梯度不会更新 pe 编码矩阵;
        4. pe 作为模型的状态被保存和加载.
        """
        # dropout 正则化来防止过拟合.
        return self.dropout(x + self.pe[:, :x.size(1)].requires_grad_(requires_grad=False))


class Embedding(torch.nn.Module):
    """
    嵌入层, d_model 输出向量的维度, vocabulary 词汇表的大小即输入的离散值的范围.
    """
    def __init__(self, *args, d_model: int, vocabulary: int, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.d_model = d_model
        self.embedding = torch.nn.Embedding(vocabulary, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 输出嵌入词向量时, 进行缩放, 让嵌入词向量数值范围更稳定, 并与位置编码数值范围相互匹配.
        return self.embedding(x) * math.sqrt(self.d_model)


class LayerNorm(torch.nn.Module):
    """
    归一化 LN 层: 函数公式见 https://s21.ax1x.com/2024/07/09/pkfrUN8.png
    """
    def __init__(self, *args, features: int, eps: float=1e-6, **kwargs) -> None:
        # features 代表归一化层特征维度大小.
        super().__init__(*args, **kwargs)
        # 创建可学习参数(训练中通过反向传播进行更新优化), γ 初始化均为 1, 不缩放.
        self.gamma = torch.nn.Parameter(torch.ones(features))
        # 创建可学习参数, β 初始化均为 0, 不偏移.
        self.beta = torch.nn.Parameter(torch.zeros(features))
        self.epsilon = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        计算 y = LN(x) 结果.
        """
        # keepdim = True, 输出与输入维度相同, 保持输出的维度.
        average = x.mean(dim=-1, keepdim=True)
        standard = x.std(dim=-1, keepdim=True)
        return self.gamma * (x - average) / torch.sqrt(standard + self.epsilon) + self.beta


class SublayerConnection(torch.nn.Module):
    """
    子层块: 残差网络连接, 模型见 https://s21.ax1x.com/2024/07/09/pkfc0I0.png

    残差网络可以在反向传播过程中, 梯度连乘, 也不会造成梯度消失.
    """
    def __init__(self, *args, size: int, dropout: float, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.norm = LayerNorm(features=size)
        # dropout 丢弃神经网络结点图示见 https://s21.ax1x.com/2024/07/09/pkfc2L9.png
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, sublayer: torch.Tensor) -> torch.Tensor:
        """
        无信息损失的 x 加上一个推理损失的函数, 这是残差网络重要思想, 避免梯度消失(求导多了常数项).
        """
        return x + self.dropout(sublayer(self.norm(x)))


class EncoderLayer(torch.nn.Module):
    """
    编码器块层: 由自"注意力"和"位置逐元素全连接前馈神经网络"构成, 多个块层可构成集群.
    """
    def __init__(
        self,
        *args,
        size: int,
        self_attention: object,
        feed_forward_network: object,
        dropout: float,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.self_attention = self_attention
        self.feed_forward_network = feed_forward_network
        # 克隆 2 块, 每个编码器层包含两个子层块.
        self.sublayer = clone(SublayerConnection(size=size, dropout=dropout), 2)
        self.size = size

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        # query, key, value 三个集合输入到 attention 函数是统一的.
        # mask 掩码指定元素 padding, 计算自注意力权重忽略遮掩的元素.
        x = self.sublayer[0](x, lambda x: self.self_attention(x, x, x, mask))
        # 把子层块的输出 x 推流至子层块的第二个, 进行 FFN 推理计算.
        return self.sublayer[1](x, self.feed_forward_network)


class DecoderLayer(torch.nn.Module):
    """
    解码器块层: 由"自注意力", "(普通)注意力", "前馈神经网络构成", 多个块层可构成集群.

    1. 解码器块比编码器块多一个(普通)注意力层, 它的输入来自解码器集群的输出 memory
    2. 自注意力和注意力相同, 只不过 query, key, value 输入不同
    3. 注意力的 query 由下层输入支持, 即自注意力的输出
    4. key, value 是编码器集群最后一层输出 memory
    5. 自注意力的 query, key, value 均由下层输入支持
    """
    def __init__(
        self,
        *args,
        size: int,
        self_attention: object,
        source_attention: object,
        feed_forward_network: object,
        dropout: float,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.size = size
        self.self_attention = self_attention
        self.source_attention = source_attention
        self.feed_forward_network = feed_forward_network
        # 克隆 3 块子层.
        self.sublayer = clone(SublayerConnection(size=size, dropout=dropout), 3)
    
    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        source_mask: object,
        target_mask: object
    ) -> torch.Tensor:
        # 自注意力.
        x = self.sublayer[0](x, lambda x: self.self_attention(x, x, x, target_mask))
        # (普通)注意力, 此处的 memory 来自于编码器的输出.
        x = self.sublayer[1](x, lambda x: self.source_attention(x, memory, memory, source_mask))
        # 前馈神经网络.
        return self.sublayer[2](x, self.feed_forward_network)


class Encoder(torch.nn.Module):
    """
    编码器集群: 编码器块堆叠, "Attention Is All You Need" 采用六层.
    """
    def __init__(self, *args, layer: object, N: int, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # 克隆拷贝 N 个编码器, 堆叠构成集群.
        self.layers = clone(layer, N)
        self.norm = LayerNorm(features=layer.size)

    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        依次通过集群的每个块层传递输入(和掩码).
        """
        # 集群中的每个块层反复推理.
        for layer in self.layers:
            x = layer(x, mask)
        # 最后输出前通过一次 LN 层.
        return self.norm(x)


class Decoder(torch.nn.Module):
    """
    解码器集群: 解码器块堆叠, 生成编码器集群将输入序列变成词向量产生的最终结果.
    """
    def __init__(self, *args, layer: object, N: int, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.layers = clone(layer, N)
        self.norm = LayerNorm(features=layer.size)
    
    def forward(
        self,
        x: torch.Tensor,
        memory: torch.Tensor,
        source_mask: object,
        target_mask: object
    ) -> torch.Tensor:
        """
        含掩码.
        """
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)


class Generator(torch.nn.Module):
    """
    生成器: 定义标准 Linear + Softmax 生成步骤.
    """
    def __init__(self, *args, d_model: int, vocabulary: int, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.projection = torch.nn.Linear(d_model, vocabulary)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.log_softmax(self.projection(x), dim=-1)


class EncoderDecoder(torch.nn.Module):
    """
    标准编码器-解码器架构, 这是未来魔改版本的基础.
    """
    def __init__(
        self,
        *args,
        encoder: object,
        decoder: object,
        source_embedding: object,
        target_embedding: object,
        generator: object,
        **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embedding
        self.tgt_embed = target_embedding
        self.generator = generator

    def encode(self, src: torch.Tensor, src_mask: torch.Tensor) -> torch.Tensor:
        return self.encoder(self.src_embed(src), src_mask)

    def decode(
        self,
        memory: torch.Tensor,
        src_mask: torch.Tensor,
        tgt: torch.Tensor,
        tgt_mask: torch.Tensor
    ) -> torch.Tensor:
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

    def forward(
        self,
        src: torch.Tensor,
        tgt: torch.Tensor,
        src_mask: torch.Tensor,
        tgt_mask: torch.Tensor
    ) -> torch.Tensor:
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)


class Transformer:
    """
    完整的 Transformer 模型, 架构见 https://s21.ax1x.com/2024/07/10/pkhVMm6.png
    """
    def __new__(
        cls,
        source_vocabulary: int,
        target_vocabulary: int,
        N: int = 6,
        d_model: int = 512,
        d_ff: int = 512,
        n_heads: int = 8,
        dropout: float = 0.1
    ) -> object:
        attn = MultiHeadAttention(
            h = n_heads,
            d_model = d_model,
            dropout = dropout
        )
        ffn = PositionWiseFeddForward(
            d_model = d_model,
            d_ff = d_ff,
            dropout = dropout
        )
        pos = PositionalEncoding(
            d_model = d_model,
            dropout = dropout,
            max_tokens = 5000
        )
        model = EncoderDecoder(
            encoder = Encoder(
                layer = EncoderLayer(
                    size = d_model,
                    self_attention = copy.deepcopy(attn),
                    feed_forward_network = copy.deepcopy(ffn),
                    dropout = dropout
                ),
                N = N
            ),
            decoder = Decoder(
                layer = DecoderLayer(
                    size = d_model,
                    self_attention = copy.deepcopy(attn),
                    source_attention = copy.deepcopy(attn),
                    feed_forward_network = copy.deepcopy(ffn),
                    dropout = dropout
                ),
                N = N
            ),
            source_embedding = torch.nn.Sequential(
                Embedding(d_model=d_model, vocabulary=source_vocabulary),
                copy.deepcopy(pos)
            ),
            target_embedding = torch.nn.Sequential(
                Embedding(d_model=d_model, vocabulary=target_vocabulary),
                copy.deepcopy(pos)
            ),
            generator = Generator(d_model=d_model, vocabulary=target_vocabulary)
        )
        # 用 Glorot / fan_avg 初始化参数.
        for param in model.parameters():
            if param.dim() > 1:
                torch.nn.init.xavier_uniform_(param)
        return model


def demoInference(epoch: int) -> None:
    """
    示例推理测试.
    """
    model = Transformer(source_vocabulary=11, target_vocabulary=11, N=1)
    model.eval()
    src = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], dtype=torch.long)
    src_mask = torch.ones(1, 1, 10)
    memory = model.encode(src, src_mask)
    prediction = torch.zeros(1, 1).type_as(src)
    for _ in range(0, 9, 1):
        output = model.decode(
            memory = memory,
            src_mask = src_mask,
            tgt = prediction,
            tgt_mask = subsequentMask(prediction.size(1)).type_as(src.data)
        )
        probability = model.generator.forward(output[:, -1])
        _, next_word = torch.max(probability, dim=1)
        next_word = next_word.data[0]
        prediction = torch.cat(
            tensors = [prediction, torch.empty(1, 1).type_as(src.data).fill_(next_word)],
            dim = 1
        )
    debug(epoch=epoch, ExampleUntrainedModel=prediction)



if __name__ == '__main__':
    for idx in range(0, 10, 1):
        demoInference(epoch = idx + 1)