Swin Transformer 变体1：使用MLP代替多头注意力机制

原创

已于 2024-12-17 13:09:28 修改 · 1.3k 阅读

23 ·

CC 4.0 BY-SA版权

文章标签：

#transformer #深度学习 #人工智能 #图像处理 #卷积神经网络 #计算机视觉

于 2024-12-13 16:19:55 首次发布

Swin Transformer v1模型具体实现：Swin Transformer模型具体代码实现-优快云博客

一、分组卷积：

分组卷积是一种卷积操作的变体，它将输入通道分成多个组，然后在每个组内独立地进行卷积操作。如果将输入的数据分为head_num个组进行卷积，分组卷积可以在不同的组内学习到不同的特征，这与多头注意力机制在不同头中捕捉不同特征的能力有相似之处。那么在图像分类领域，可以使用MLP代替多头注意力机制，这样做的好处有：

减少参数和计算量：通过在每个组内进行卷积，分组卷积显著减少了模型的参数和计算量。
特征多样性：允许模型在不同的通道组中学习不同的特征，增加了模型的表达能力。
计算效率：分组卷积在计算上比多头注意力机制更高效，尤其是在处理高维特征时。多头注意力机制的计算复杂度通常是O( $n^{2}$ )O( $n^{2}$ )，而分组卷积的复杂度更低，适合于大规模数据处理。
简化模型结构：使用分组卷积可以简化模型结构，减少模型复杂度，同时保持或提高性能。这对于需要快速推理的应用场景（如实时处理）非常重要。

二、具体实现：

1、替换注意力模块完整代码：

class SwinMLPBlock(nn.Module):
    def __init__(self, 
                 dim, 
                 input_resolution, #分辨率
                 num_heads, 
                 window_size=7, 
                 shift_size=0,
                 mlp_ratio=4., 
                 drop=0., 
                 drop_path=0.,
                 act_layer=nn.GELU, 
                 norm_layer=nn.LayerNorm):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_si