从零开始实现 MobileViT 注意力机制——轻量级Transformer Vision Model 的新思路

从零开始实现 MobileViT 注意力机制——轻量级Transformer Vision Model 的新思路

近年来，计算机视觉领域中 Transformer 模型的崛起为图像处理带来了新的活力。特别是在 ViT（Vision Transformer）模型提出之后，Transformer 在图像分类、目标检测等任务上展示了超越 CNN 的潜力。然而，标准的 ViT 模型参数量大，计算复杂度高，难以在移动设备等资源受限的环境中部署。

最近，《MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer》 这篇论文提出了一种轻量化、通用且适合移动端的视觉变换器模型。该模型通过结合局部和全局特征的创新设计，在保持良好性能的同时，大大降低了计算资源的需求，为移动应用提供了新的解决方案。

本文将从零开始解读并实现 MobileViT 的核心注意力机制模块，帮助开发者理解这一轻量级视觉变换器的工作原理，从而在实际项目中灵活运用。

---

1. 背景：从 ViT 到 MobileViT

1.1 Vision Transformer (ViT) 简介

标准的 ViT 模型将整个图像划分为不重叠的 patches（块），并将其转换为序列输入到基于Transformer 的编码器中。这种方法虽然在性能上表现出色，但也带来了以下问题：

• 计算复杂度高：将图像分割成大量 patches 后进行序列操作，参数量和计算量急剧上升。
• 适用性有限：直接使用 Transformer 架构处理图像分辨率较高的场景时，资源消耗（如内存、算力）难以满足移动端的需求。

1.2 MobileViT 的创新思路

MobileViT 提出了一种折中的解决方案——结合 局部表示（Local Representation） 和 全局表示（Global Representation），以降低计算复杂度同时保持性能。其核心思想是：

• 在每个位置保留原始图像的局部特征信息。
• 通过 Transformer 模块提取和增强全局特征信息。
• 将局部和全局特征进行融合，生成最终的高质量视觉表征。

---

2. MobileViT 注意力机制模块实现解析

MobileViT 的核心模块是 MobileViTAttention。我们需要逐步解读其实现细节，并通过代码示例帮助读者理解其工作原理。

2.1 模块设计概述

• 输入：一个张量（Tensor），形状为 [batch_size, in_channel, height, width]
• 输出：经过局部和全局特征融合后的张量，保持与输入相同的尺寸

模块主要包含以下几个部分：

1. 局部特征提取：通过卷积操作提取每个位置的局部信息。
2. 全局特征提取：使用 Transformer 模块对图像进行分块（patch）处理，并在序列空间中捕获长距离依赖关系。
3. 特征融合：将局部和全局特征拼接后，通过轻量级的卷积操作生成最终输出。

以下是完整的 MobileViTAttention 类的实现代码：

import torch
from torch import nn

class MobileViT_Attention(nn.Module):
    def __init__(self, in_channels=3, kernel_size=3, patch_size=2, embed_dim=144):
        super().__init__()
        # 设置 patch 的大小（默认为7x7）
        self.ph, self.pw = patch_size, patch_size
      
        # 局部特征提取：通过卷积操作捕获局部上下文信息
        self.local_conv = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, 
                      padding=kernel_size//2, stride=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True)
        )
      
        # 全局特征提取：将张量重排为 [batch_size, patch_height*patch_width, N_h*N_w, embed_dim]
        # Transformer 模块用于捕获全局上下文信息
        self.global_trans = Transformer(embed_dim=embed_dim,
                                        num_heads=16,
                                        num_transformer_layers=4)
      
        # 特征融合：将局部特征和全局特征拼接，并通过卷积操作生成最终输出
        self.fusion_conv = nn.Sequential(
            nn.Conv2d(in_channels*2, in_channels, kernel_size=kernel_size, 
                      padding=kernel_size//2, stride=1),
            nn.BatchNorm2d(in_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        # 提取局部特征
        local_feats = self.local_conv(x)  # 局部特征

        if len(local_feats.shape) == 4:
            B, C, H, W = local_feats.shape
        else:
            raise ValueError("Input tensor should have rank 4.")

        # 分割图像为 patch，并进行重排：从 [B, C, H, W] 到 [B, (H*W), C]
        # 每个 patch 的大小为 (patch_size, patch_size)
        patches = []
        for i in range(0, H, self.ph):
            for j in range(0, W, self.pw):
                patch = local_feats[:, :, i:i+self.ph, j:j+self.pw]
                patch = torch.flatten(patch, start_dim=2)  # 打平patch
                patches.append(patch)
      
        # 拼接所有的 patch，形成张量 [B, num_patches, C]
        x_patched = torch.stack(patches, dim=1)

        # 传递到 Transformer 中提取全局特征
        global_feats = self.global_trans(x_patched)  # 全局上下文特征

        # 特征融合：将原始输入的局部特征与 Transformer 输出的全局特征拼接
        x_fused = torch.cat([local_feats, global_feats.unsqueeze(2).unsqueeze(3)], dim=1)

        return self.fusion_conv(x_fused)  # 最终的特征输出

class Transformer(nn.Module):
    def __init__(self, embed_dim=768, num_heads=12, 
                 num_transformer_layers=4):
        super().__init__()
        self.embedding = nn.Linear(embed_dim, embed_dim)
        self.layers =(nn.ModuleList([
            TransformerBlock(d_model=embed_dim, nhead=num_heads)
            for _ in range(num_transformer_layers)
        ]))
      
    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        return x

---

3. 实现细节解读

3.1 局部特征提取

• 卷积操作：使用 nn.Conv2d 在局部区域内捕获上下文信息。
• BN 和 ReLU：通过归一化和非线性激活，提升特征表达能力。

self.local_conv = nn.Sequential(
    nn.Conv2d(3, 3, kernel_size=3, padding=1),
    nn.BatchNorm2d(3),
    nn.ReLU(inplace=True)
)

3.2 全局特征提取（Transformer）

• 分块：将图像分割为 patch_size x patch_size 的小块，每个块展开成一维向量。
• 序列建模：通过多层 Transformer Block 捕获长距离依赖。

class TransformerBlock(nn.Module):
    def __init__(self, d_model=768, nhead=12):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(embed_dim=d_model, num_heads=nhead)
        self.dropout = nn.Dropout(0.1)
      
    def forward(self, x):
        out = self.self_attn(x, x, x)[0]
        return F.dropout(out, p=0.1, training=self.training)

3.3 特征融合

• 拼接：将局部特征和全局特征在通道维度上进行拼接。
• 卷积操作：通过轻量级的卷积操作生成最终输出。

self.fusion_conv = nn.Sequential(
    nn.Conv2d(3*2, 3, kernel_size=3, padding=1),
    nn.BatchNorm2d(3),
   nn.ReLU(inplace=True)
)

---

4. 模块的输入输出尺寸分析

输入

• 形状：[batch_size, in_channels, height, width]
• 示例：[ batch_size: 4, in_channels: 3 (RGB), height: 224, width: 224 ]

输出

• 相同的尺寸 [batch_size, in_channels, height, width]
• 经过局部和全局特征融合后，输出高质量的视觉表征。

---

5. 总结与展望

通过结合局部和全局特征提取，MobileViT 成功地在轻量级计算资源的基础上实现了高效的视觉信息处理。这一模块尤其适合应用于移动设备和嵌入式系统中，同时也可以作为其他视觉任务（如目标检测、图像分割）的高效特征提取模块。

未来的工作可以尝试以下方向：

1. 优化 Transformer 模块：通过减少头数或简化注意力机制降低计算复杂度。
2. 自适应 patch 大小：根据输入尺寸动态调整 patch 的大小，提升模型的灵活性。
3. 多尺度融合：在更细粒度的尺度上结合特征信息，进一步提升性能。

希望通过对这一模块的解读和实现，能够帮助读者更好地理解和应用 MobileViT 模型，在实际项目中发挥其优势。