【DeepLearning-5】基于Transformer架构的自定义神经网络类

类定义

class UserDefined(nn.Module):

• UserDefined 是一个自定义的神经网络类，继承自PyTorch的 nn.Module 基类。

构造函数 __init__

def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):

• __init__ 方法是类的构造函数，用于初始化 UserDefined 实例。
• dim: 特征的维度。
• depth: Transformer的层数。
• heads: 多头注意力机制中的头数。
• dim_head: 每个头的维度。
• mlp_dim: 前馈网络中隐藏层的维度。
• dropout: 用于正则化的dropout率。

Transformer层的初始化

self.layers = nn.ModuleList([]) 
for _ in range(depth): 
    self.layers.append(nn.ModuleList([ 
        PreNorm(dim, Attention(dim, heads, dim_head, dropout)), 
        PreNorm(dim, FeedForward(dim, mlp_dim, dropout)) ]))

• 使用 nn.ModuleList 创建一个空的模块列表，用于存储Transformer的层。

"Transformer的层"是指构成Transformer架构中的基本单元或模块。在Transformer架构中，整个模型是由多个这样的层堆叠而成的。每一层通常包含以下两个主要组件：

1. 多头注意力机制（Multi-head Attention）: 这是Transformer模型的核心特征之一。它允许模型在处理一个序列时同时关注序列的不同部分。这种机制通过将注意力分散到多个头上来实现，每个头学习序列的不同方面。

2. 前馈神经网络（Feed Forward Neural Network）: 这是一个简单的神经网络，通常包含两个线性变换和一个激活函数。在Transformer层中，它跟在多头注意力机制之后，进一步处理数据。

UserDefined 类中，每个Transformer层由以下两部分组成：

• 带有归一化的多头注意力（PreNorm(dim, Attention(dim, heads, dim_head, dropout))）: 这里的 PreNorm 是一个预归一化层，用于在多头注意力机制之前对输入进行归一化处理。Attention 是实现多头注意力机制的部分。

• 带有归一化的前馈网络（PreNorm(dim, FeedForward(dim, mlp_dim, dropout))）: 这同样使用 PreNorm 进行预归一化，然后通过 FeedForward 层进行进一步的数据处理。

在 for 循环中创建了多个Transformer层，并将它们添加到 self.layers（详见PyTorch中self.layers的作用-CSDN博客） 这个 nn.ModuleList 中。使用 nn.ModuleList 是因为它可以正确地注册其内部的模块，以便PyTorch能够跟踪它们的参数，使得这些参数可以在训练过程中被优化。

• for _ in range(depth) 循环 depth 次，每次循环中：
  • 创建一个 PreNorm 层，包装一个 Attention 层。Attention 层实现了多头注意力机制。
  • 创建另一个 PreNorm 层，包装一个 FeedForward 层。FeedForward 是一个简单的前馈网络。
  • 这两个层被添加到内部的 nn.ModuleList 中，形成一个Transformer的编码层。

前向传播 forward

def forward(self, x): 
    for attn, ff in self.layers: 
        x = attn(x) + x 
        x = ff(x) + x 
    return x

• forward 方法定义了数据通过网络的方式。
• 对于每个Transformer层：
  • 应用注意力层 (attn) 和残差连接，更新 x。
  • 接着应用前馈层 (ff) 和另一个残差连接，再次更新 x。

完整代码：

class UserDefined(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
            ]))
    
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x