Vision Transformer实战指南：从模型选型到部署优化的全流程解析

1次阅读

共计 2090 个字符，预计需要花费 6 分钟才能阅读完成。

传统 CNN（Convolutional Neural Networks）通过局部感受野逐步提取特征，但在处理长距离依赖（如遥感图像中的跨区域关联）时存在天然局限。ViT（Vision Transformer）的突破性在于：

全局建模能力：通过 Self-Attention 机制直接建立图像所有区域的关联
参数效率：在大规模数据上，相同参数量下 ViT 常优于 CNN
架构统一性：与 NLP 领域 Transformer 共享基础组件，便于多模态建模

模型	计算复杂度 (FLOPs)	ImageNet Top-1 Acc	训练显存占用
ResNet-50	4.1G	76.1%	7.8GB
EfficientNet-B3	1.8G	81.1%	6.2GB
ViT-Base/16	17.6G	84.5%	12.4GB

注：测试环境为 NVIDIA V100，batch_size=256

将图像分割为 $P \times P$ 的块，每个块展平为 $P^2 \times C$ 的向量（$C$ 为通道数）。数学表达：

$$
z_0 = [x_{class}; x_p^1E; x_p^2E; …; x_p^NE] + E_{pos}
$$

其中 $E \in \mathbb{R}^{(P^2 \cdot C) \times D}$ 为可学习的投影矩阵，$E_{pos}$ 为位置编码。PyTorch 实现：

# 要求 torch>=1.9.0
class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, embed_dim, 
                             kernel_size=patch_size, 
                             stride=patch_size)  # 用卷积实现分块投影

    def forward(self, x):
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

额外添加的可学习向量，最终输出作为整个图像的表示
与 NLP 中的[CLS] token 类似，通过 Attention 聚合全局信息

def scaled_dot_product_attention(q, k, v, mask=None):
    # 输入张量形状: (batch, heads, seq_len, dim)
    matmul_qk = torch.matmul(q, k.transpose(-2, -1))
    dk = q.size()[-1]
    scaled_attention_logits = matmul_qk / math.sqrt(dk)

    if mask is not None:
        scaled_attention_logits += (mask * -1e9)  

    attention_weights = F.softmax(scaled_attention_logits, dim=-1)
    output = torch.matmul(attention_weights, v)
    return output

# 启用 CUDA 原生 Attention（PyTorch 1.12+）torch.backends.cuda.enable_flash_sdp(True)

减小 Layer Norm 的 eps 值（如从 1e- 6 调至 1e-8）
使用 Pre-LN 结构代替 Post-LN
添加 MixUp/CutMix 数据增强

scaler = torch.cuda.amp.GradScaler()  # 必须使用梯度缩放

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()