【扒代码】transformer.py

from .mlp import MLP  # 从当前包中导入MLP模块

from torch import nn  # 从PyTorch库中导入nn模块

# 定义TransformerEncoder类，继承自nn.Module
class TransformerEncoder(nn.Module):

    def __init__(  # 初始化函数，定义类的初始化参数
        self,
        num_layers: int,  # Transformer Encoder的层数
        emb_dim: int,  # 嵌入维度
        num_heads: int,  # 多头注意力机制的头数
        dropout: float,  # dropout概率
        layer_norm_eps: float,  # LayerNorm的epsilon值
        mlp_factor: int,  # MLP因子
        norm_first: bool,  # 是否先进行规范化
        activation: nn.Module,  # 激活函数
        norm: bool,  # 是否进行最终的规范化
    ):

        super(TransformerEncoder, self).__init__()  # 调用父类的初始化函数

        # 使用列表推导式创建一个ModuleList，其中包含num_layers个TransformerEncoderLayer实例
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(
                emb_dim, num_heads, dropout, layer_norm_eps,
                mlp_factor, norm_first, activation
            ) for _ in range(num_layers)
        ])

        # 如果norm为真，定义一个LayerNorm层，否则定义一个Identity层
        self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity()

    # 定义前向传播函数
    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
        output = src  # 将输入赋值给output变量
        for layer in self.layers:  # 遍历每一层
            # 将output、pos_emb、src_mask、src_key_padding_mask传入每一层
            output = layer(output, pos_emb, src_mask, src_key_padding_mask)
        # 返回经过所有层和最后规范化处理后的结果
        return self.norm(output)

功能解释

类和方法

• TransformerEncoder类: 这个类实现了一个Transformer编码器。
• 初始化函数__init__: 初始化编码器的各个参数，并创建指定层数的Transformer编码器层。
• 前向传播函数forward: 定义前向传播过程，将输入依次传递给各个编码器层，并进行最终的规范化处理。

参数

• num_layers: Transformer编码器的层数。
• emb_dim: 输入嵌入的维度。
• num_heads: 多头注意力机制的头数。
• dropout: Dropout层的概率，用于防止过拟合。
• layer_norm_eps: LayerNorm层中的epsilon值，防止分母为零。
• mlp_factor: MLP的扩展因子，用于确定MLP的隐藏层大小。
• norm_first: 一个布尔值，指示是否先进行规范化。
• activation: 激活函数，用于引入非线性。
• norm: 一个布尔值，指示是否在输出前进行规范化。

方法

• __init__方法: 初始化Transformer编码器层和规范化层。
• forward方法: 实现前向传播，依次通过每一层Transformer编码器层，并在最后进行规范化处理。

import torch
from torch import nn

class TransformerEncoderLayer(nn.Module):
    def __init__(
        self,
        emb_dim: int,              # 嵌入维度
        num_heads: int,            # 注意力机制中的头数
        dropout: float,            # dropout比率
        layer_norm_eps: float,    # 层归一化中的epsilon值
        mlp_factor: int,          # MLP的扩展因子
        norm_first: bool,          # 是否先进行归一化
        activation: nn.Module,     # 激活函数模块
    ):
        super(TransformerEncoderLayer, self).__init__()  # 调用基类的初始化方法

        # 是否先进行归一化
        self.norm_first = norm_first

        # 创建层归一化层
        self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps)
        self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps)
        # 创建dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        # 创建自注意力机制
        self.self_attn = nn.MultiheadAttention(
            emb_dim, num_heads, dropout
        )
        # 创建多层感知机
        self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation)

    # 辅助函数，用于将输入x与嵌入emb结合
    def with_emb(self, x, emb):
        return x if emb is None else x + emb

    # 前向传播方法
    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
        # 如果norm_first为True，则先进行归一化再进行自注意力和MLP操作
        if self.norm_first:
            src_norm = self.norm1(src)
            # 查询（q）、键（k）和值（v）使用相同的归一化后的输入加上位置编码
            q = k = src_norm + pos_emb
            # 进行自注意力操作
            src = src + self.dropout1(self.self_attn(
                query=q,
                key=k,
                value=src_norm,
                attn_mask=src_mask,
                key_padding_mask=src_key_padding_mask
            )[0])

            src_norm = self.norm2(src)
            # 通过MLP处理归一化后的输入
            src = src + self.dropout2(self.mlp(src_norm))

        # 如果norm_first为False，则先进行自注意力操作再进行归一化和MLP操作
        else:
            q = k = src + pos_emb
            src = self.norm1(src + self.dropout1(self.self_attn(
                query=q,
                key=k,
                value=src,
                attn_mask=src_mask,
                key_padding_mask=src_key_padding_mask
            )[0]))

            src = self.norm2(src + self.dropout2(self.mlp(src)))

        # 返回编码器层的输出
        return src

这段代码定义了一个名为 TransformerEncoderLayer 的类，它是 PyTorch nn.Module 的子类，用于实现 Transformer 编码器层。编码器层是 Transformer 模型中的基本构建块，负责处理输入序列并提取特征。

怎么只有编码器？

功能解释：

• TransformerEncoderLayer 类接收多个参数来配置 Transformer 编码器层的行为，包括嵌入维度、注意力头数、dropout 比率等。
• self_attn 是自注意力机制，它允许模型在处理序列时考虑序列内部的依赖关系。
• mlp 是多层感知机，用于在注意力机制之后进一步处理特征。
• with_emb 函数用于将输入特征 x 与位置编码 emb 结合，如果位置编码为 None，则直接返回输入特征。
• forward 方法定义了编码器层的前向传播逻辑，它接收源序列 src、位置编码 pos_emb、源掩码 src_mask 和源键填充掩码 src_key_padding_mask 作为输入。
• 根据 norm_first 参数的值，编码器层可以采用不同的处理流程，即先进行归一化再进行注意力和 MLP 操作，或者先进行注意力操作再进行归一化和 MLP 操作。
• 方法返回编码器层处理后的特征表示。

整体而言，TransformerEncoderLayer 类实现了 Transformer 模型中的一个编码器层，它可以捕捉序列数据的内部依赖关系，并通过 MLP 对特征进行非线性变换，以提取丰富的特征表示。这种设计使得 Transformer 模型在处理序列数据时非常灵活和强大。

完整：

from .mlp import MLP

from torch import nn


# 定义TransformerEncoder类，继承自nn.Module
class TransformerEncoder(nn.Module):

    def __init__(
        self,
        num_layers: int,# Transformer Encoder的层数
        emb_dim: int,# 嵌入维度
        num_heads: int,# 多头注意力机制的头数
        dropout: float,# dropout概率
        layer_norm_eps: float,# LayerNorm的epsilon值
        mlp_factor: int,# MLP因子
        norm_first: bool,# 是否先进行规范化
        activation: nn.Module,# 激活函数
        norm: bool,# 是否进行最终的规范化
    ):

        super(TransformerEncoder, self).__init__()# 调用父类的初始化函数

        # 使用列表推导式创建一个ModuleList，其中包含num_layers个TransformerEncoderLayer实例
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(
                emb_dim, num_heads, dropout, layer_norm_eps,
                mlp_factor, norm_first, activation
            ) for _ in range(num_layers)
        ])

        # 如果norm为真，定义一个LayerNorm层，否则定义一个Identity层
        self.norm = nn.LayerNorm(emb_dim, layer_norm_eps) if norm else nn.Identity()

    # 定义前向传播函数
    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
        output = src # 将输入赋值给output变量
        for layer in self.layers: # 遍历每一层
            # 将output、pos_emb、src_mask、src_key_padding_mask传入每一层
            output = layer(output, pos_emb, src_mask, src_key_padding_mask)
        # 返回经过所有层和最后规范化处理后的结果
        return self.norm(output)


# 实现 Transformer 编码器层。编码器层是 Transformer 模型中的基本构建块，负责处理输入序列并提取特征
class TransformerEncoderLayer(nn.Module):

    def __init__(
        self,
        emb_dim: int,# 嵌入维度
        num_heads: int,# 注意力机制中的头数
        dropout: float,# dropout比率
        layer_norm_eps: float,# 层归一化中的epsilon值
        mlp_factor: int,# MLP的扩展因子
        norm_first: bool,# 是否先进行归一化
        activation: nn.Module,# 激活函数模块
    ):
        super(TransformerEncoderLayer, self).__init__()# 调用基类的初始化方法

        # 是否先进行归一化
        self.norm_first = norm_first

        # 创建层归一化层
        self.norm1 = nn.LayerNorm(emb_dim, layer_norm_eps)
        self.norm2 = nn.LayerNorm(emb_dim, layer_norm_eps)
        # 创建dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        # 创建自注意力机制
        self.self_attn = nn.MultiheadAttention(
            emb_dim, num_heads, dropout
        )
        # 创建多层感知机
        self.mlp = MLP(emb_dim, mlp_factor * emb_dim, dropout, activation)

    # 辅助函数，用于将输入x与嵌入emb结合
    def with_emb(self, x, emb):
        return x if emb is None else x + emb

    # 前向传播方法
    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
        # 如果norm_first为True，则先进行归一化再进行自注意力和MLP操作
        if self.norm_first:
            src_norm = self.norm1(src)
            # 查询（q）、键（k）和值（v）使用相同的归一化后的输入加上位置编码
            q = k = src_norm + pos_emb
            # 进行自注意力操作
            src = src + self.dropout1(self.self_attn(
                query=q,
                key=k,
                value=src_norm,
                attn_mask=src_mask,
                key_padding_mask=src_key_padding_mask
            )[0])

            src_norm = self.norm2(src)
             # 通过MLP处理归一化后的输入
            src = src + self.dropout2(self.mlp(src_norm))
        # 如果norm_first为False，则先进行自注意力操作再进行归一化和MLP操作
        else:
            q = k = src + pos_emb
            src = self.norm1(src + self.dropout1(self.self_attn(
                query=q,
                key=k,
                value=src,
                attn_mask=src_mask,
                key_padding_mask=src_key_padding_mask
            )[0]))

            src = self.norm2(src + self.dropout2(self.mlp(src)))

        # 返回编码器层的输出
        return src