【Block总结】门控结构的MLP结构

AI浩

发布于 2025-01-16 10:29:16

9400

代码可运行

文章被收录于专栏：AI智韵

运行总次数：0

代码可运行

模块

记录一个具有门控模块的MLP，这个模块可以降低MLP的参数量，还可以提高模型的精度，很多模型都用到了这样的结构，代码如下：

class Gate(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.conv = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, groups=dim) # DW Conv

    def forward(self, x, H, W):
        # Split
        x1, x2 = x.chunk(2, dim = -1)
        B, N, C = x.shape
        x2 = self.conv(self.norm(x2).transpose(1, 2).contiguous().view(B, C//2, H, W)).flatten(2).transpose(-1, -2).contiguous()

        return x1 * x2


class MLP(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.sg = Gate(hidden_features//2)
        self.fc2 = nn.Linear(hidden_features//2, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x, H, W):
        """
        Input: x: (B, H*W, C), H, W
        Output: x: (B, H*W, C)
        """
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)

        x = self.sg(x, H, W)
        x = self.drop(x)

        x = self.fc2(x)
        x = self.drop(x)
        return x

这个代码定义了两个类：Gate和MLP，它们都是基于PyTorch框架实现的神经网络模块。下面是对这两个类及其功能的讲解：

Gate 类

Gate类是一个自定义的神经网络层，主要用于对输入特征进行特定的变换。它的结构如下：

初始化 (__init__ 方法):
- dim: 输入特征的维度。
- self.norm: 一个Layer Normalization层，用于对输入特征进行归一化处理，有助于加速训练过程并提高模型的稳定性。
- self.conv: 一个深度可分离卷积层（Depthwise Convolution，通过设置groups=dim实现），卷积核大小为3x3，步长为1，填充为1。这意味着卷积操作在每个输入通道上独立进行，有助于捕捉局部特征。
前向传播 (forward 方法):
- 输入x的形状为(B, N, C)，其中B是批次大小，N是特征的数量（可能是空间维度H*W的展平），C是特征维度。
- x被沿着最后一个维度分成两部分x1和x2。
- x2首先经过Layer Normalization，然后重塑并转置以适配卷积层的输入要求，接着进行深度可分离卷积，最后再次重塑和转置以恢复原始形状的一部分维度。
- 输出是x1与变换后的x2的逐元素乘积，这种操作可能有助于特征之间的交互和信息流动。

MLP 类

MLP类是一个多层感知机（Multilayer Perceptron），其结构如下：

初始化 (__init__ 方法):
- in_features: 输入特征的维度。
- hidden_features: 隐藏层的特征维度，默认为输入特征的维度。
- out_features: 输出特征的维度，默认为输入特征的维度。
- act_layer: 激活函数层，默认为GELU（Gaussian Error Linear Unit）。
- drop: Dropout比率，用于减少过拟合。
- self.fc1: 第一个全连接层，将输入特征映射到隐藏层特征空间。
- self.act: 激活函数层。
- self.sg: 一个Gate层，对隐藏层特征的一部分进行特定的变换。
- self.fc2: 第二个全连接层，将变换后的隐藏层特征映射到输出特征空间。
- self.drop: Dropout层。
前向传播 (forward 方法):
- 输入x的形状为(B, H*W, C)，其中B是批次大小，H*W是空间维度的展平，C是特征维度。
- 输入x首先经过第一个全连接层、激活函数层和Dropout层。
- 然后，x被传递给Gate层进行处理，这一步可能涉及特征的重新组合和局部信息的捕捉。
- 经过Gate层处理后，x再次经过Dropout层，然后传递给第二个全连接层。
- 最后，输出经过另一个Dropout层处理，得到最终的输出。