【论文+代码】VISION PERMUTATOR 即插即用的多层感知器（MLP）模块

论文

本文的研究成果在项目的实现过程中起到了至关重要的作用。以下是本文的详细信息：

文章链接: VISION PERMUTATOR: A PERMUTABLE MLP-LIKE ARCHITECTURE FOR VISUAL RECOGNITION

模块

创新点

在多个方面进行了创新和改进，以下是项目的主要创新点：

代码

代码链接 https://github.com/houqb/VisionPermutator/blob/main

模块

核心算法和模型训练。

import torch.nn as nn

class VisionPermutator(nn.Module):
    """ Vision Permutator
    视觉排列器，是一个用于图像处理的神经网络模型
    """
    def __init__(self, layers, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
        embed_dims=None, transitions=None, segment_dim=None, mlp_ratios=None, skip_lam=1.0,
        qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
        norm_layer=nn.LayerNorm, mlp_fn=WeightedPermuteMLP):

        super().__init__()
        self.num_classes = num_classes

        # 将输入图像分割成小块，并进行初步的嵌入
        self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dims[0])

        network = []
        # 构建整个网络的架构
        for i in range(len(layers)):
            # 添加基础模块到网络中
            stage = basic_blocks(embed_dims[i], i, layers, segment_dim[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias,
                    qk_scale=qk_scale, attn_drop=attn_drop_rate, drop_path_rate=drop_path_rate, norm_layer=norm_layer, skip_lam=skip_lam,
                    mlp_fn=mlp_fn)
            network.append(stage)
            if i >= len(layers) - 1:
                break
            if transitions[i] or embed_dims[i] != embed_dims[i+1]:
                # 如果需要转换，添加下采样层
                patch_size = 2 if transitions[i] else 1
                network.append(Downsample(embed_dims[i], embed_dims[i+1], patch_size))

        # 将网络模块列表封装成 nn.ModuleList
        self.network = nn.ModuleList(network)

        # 添加归一化层
        self.norm = norm_layer(embed_dims[-1])

        # 分类头部，进行最终的类别预测
        self.head = nn.Linear(embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()
        self.apply(self._init_weights)

    def _init_weights(self, m):
        # 初始化模型权重
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_classifier(self):
        # 获取分类器
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        # 重置分类器，用于迁移学习或微调
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_embeddings(self, x):
        # 处理输入图像，进行初步的嵌入
        x = self.patch_embed(x)
        # 将张量维度从 (B, C, H, W) 转换为 (B, H, W, C)
        x = x.permute(0, 2, 3, 1)
        return x

    def forward_tokens(self, x):
        # 依次通过网络的每个模块
        for idx, block in enumerate(self.network):
            x = block(x)
        # 展平张量以适应分类头部的输入要求
        B, H, W, C = x.shape
        x = x.reshape(B, -1, C)
        return x

    def forward(self, x):
        # 前向传播的入口函数
        # 首先，进行初步嵌入处理
        x = self.forward_embeddings(x)
        # 接着，通过所有网络模块进行处理
        x = self.forward_tokens(x)
        # 最后，进行归一化处理，并通过分类头部进行类别预测
        x = self.norm(x)
        # 将所有位置的特征平均，得到最终的预测结果
        return self.head(x.mean(1))

分析

https://github.com/houqb/VisionPermutator/blob/main/models/vip.py

import torch.nn as nn

class VisionPermutator(nn.Module):
    """ Vision Permutator
    视觉排列器，是一个用于图像处理的神经网络模型
    """
    def __init__(self, layers, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
        embed_dims=None, transitions=None, segment_dim=None, mlp_ratios=None, skip_lam=1.0,
        qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
        norm_layer=nn.LayerNorm, mlp_fn=WeightedPermuteMLP):

        super().__init__()
        self.num_classes = num_classes

        # 将输入图像分割成小块，并进行初步的嵌入
        self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dims[0])

        network = []
        # 构建整个网络的架构
        for i in range(len(layers)):
            # 添加基础模块到网络中
            stage = basic_blocks(embed_dims[i], i, layers, segment_dim[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias,
                    qk_scale=qk_scale, attn_drop=attn_drop_rate, drop_path_rate=drop_path_rate, norm_layer=norm_layer, skip_lam=skip_lam,
                    mlp_fn=mlp_fn)
            network.append(stage)
            if i >= len(layers) - 1:
                break
            if transitions[i] or embed_dims[i] != embed_dims[i+1]:
                # 如果需要转换，添加下采样层
                patch_size = 2 if transitions[i] else 1
                network.append(Downsample(embed_dims[i], embed_dims[i+1], patch_size))

        # 将网络模块列表封装成 nn.ModuleList
        self.network = nn.ModuleList(network)

        # 添加归一化层
        self.norm = norm_layer(embed_dims[-1])

        # 分类头部，进行最终的类别预测
        self.head = nn.Linear(embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()
        self.apply(self._init_weights)

    def _init_weights(self, m):
        # 初始化模型权重
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_classifier(self):
        # 获取分类器
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        # 重置分类器，用于迁移学习或微调
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def forward_embeddings(self, x):
        # 处理输入图像，进行初步的嵌入
        x = self.patch_embed(x)
        # 将张量维度从 (B, C, H, W) 转换为 (B, H, W, C)
        x = x.permute(0, 2, 3, 1)
        return x

    def forward_tokens(self, x):
        # 依次通过网络的每个模块
        for idx, block in enumerate(self.network):
            x = block(x)
        # 展平张量以适应分类头部的输入要求
        B, H, W, C = x.shape
        x = x.reshape(B, -1, C)
        return x

    def forward(self, x):
        # 前向传播的入口函数
        # 首先，进行初步嵌入处理
        x = self.forward_embeddings(x)
        # 接着，通过所有网络模块进行处理
        x = self.forward_tokens(x)
        # 最后，进行归一化处理，并通过分类头部进行类别预测
        x = self.norm(x)
        # 将所有位置的特征平均，得到最终的预测结果
        return self.head(x.mean(1))

代码讲解

forward 函数是模型的主要执行路径，它将输入图像转化为嵌入特征，通过网络模块进行处理，归一化后通过分类头部输出预测结果。这一步步的处理确保了模型能够有效地提取图像特征，并进行准确的分类。

1. 初始化函数 (__init__)：

   • 初始化模型，设置各种参数，并构建网络结构。
   • patch_embed 将输入图像分割成小块并进行初步的嵌入。
   • network 是一个包含多个阶段和下采样层的模块列表。
   • norm 和 head 分别是归一化层和分类头部，用于最终的预测。

2. 权重初始化函数 (_init_weights)：

   • 初始化线性层和归一化层的权重和偏置。

3. 获取分类器函数 (get_classifier)：

   • 返回分类头部，主要用于评估或微调。

4. 重置分类器函数 (reset_classifier)：

   • 重置分类头部，以适应不同的分类任务。

5. 前向嵌入函数 (forward_embeddings)：

   • 对输入图像进行初步处理，将其分割成小块并嵌入到更高维度的特征空间。

6. 前向处理函数 (forward_tokens)：

   • 通过所有网络模块进行处理，将特征进一步提取和转换。

7. 前向传播函数 (forward)：

   • 这个函数是整个模型的核心：
     1. 首先调用 forward_embeddings 对输入图像进行初步处理，将其转换为嵌入特征。
     2. 然后调用 forward_tokens 依次通过所有网络模块进行特征提取和转换。
     3. 接着对特征进行归一化处理。
     4. 最后，通过分类头部进行类别预测。这里使用了全局平均池化 (x.mean(1))，将所有位置的特征平均，得到一个最终的特征向量，输入到分类头部进行预测。