Vision Transformer (ViT)详解：从图像到序列的视觉革命

引言

2020年，Google发布的《An Image is Worth 16x16 Words》论文彻底打破了CNN在计算机视觉领域的垄断格局。Vision Transformer (ViT)首次成功将NLP界的“万能模型”Transformer迁移到图像分类，并在大规模数据集上实现了对ResNet等经典CNN的性能反超。

ViT的核心创新只有一句话，但意义深远：

像处理自然语言句子一样处理图像——把图像切割成“视觉词汇”的序列，用自注意力机制直接捕捉全局依赖。

1. ViT的诞生：CNN的“天花板”与破局

1.1 为什么CNN不够用了？

尽管ResNet、EfficientNet等CNN把局部建模做到了极致，但它们天生带着两个难以突破的归纳偏置：

归纳偏置带来的优势隐含的局限性
局部感受野学习边缘、纹理等基础特征快需堆叠几十层才能勉强“看到”全局（如猫狗的身体轮廓）
平移不变性对物体位置变化鲁棒缺乏对位置关系的显式建模（如“猫脸在猫耳朵下方”）
静态权重推理效率高所有图像用同一套卷积核，无法动态关注“当前图像的关键区域”

1.2 ViT的破局思路

ViT直接把CNN的“局部优先”推倒重来，用Transformer的全局优先范式重新定义视觉建模：

  1. ✅ 全局感知一步到位：第一层自注意力就能让图像左上角的像素“看见”右下角的像素
  2. ✅ 注意力权重动态生成：根据输入图像内容，自动调整不同区域的重要性
  3. ✅ 可扩展性极强：模型越大、数据越多，性能提升越明显（ scaling law 在视觉领域的完美验证）

2. 极简架构拆解：ViT到底做了哪几件事？

ViT的主体几乎完全照搬了NLP Transformer的编码器，唯一的改动是把“文本序列”换成了“视觉序列”。核心流程可以用4个步骤概括：

graph TD
    A[输入图像<br/>224×224×3] --> B[切割为不重叠的块<br/>16×16×3 × 196块]
    B --> C[线性投影<br/>每块→768维向量]
    C --> D[拼接CLS Token + 位置编码<br/>197×768]
    D --> E[堆叠N层Transformer编码器]
    E --> F[取CLS Token输出<br/>分类头→1000类]

关键组件速览

  1. Patch Embedding（图像转序列的核心）：用卷积或展平+线性层，把图像块转成固定维度的向量
  2. CLS Token：学习一个“全局汇总向量”，最终用它做分类（借鉴BERT）
  3. 可学习位置编码：把图像块的“位置信息”注入向量（因为Transformer本身是“位置无关”的）
  4. Transformer编码器：多头自注意力 + MLP + 残差连接 + 层归一化

3. PyTorch极简实现：从零搭ViT-B/16

我们来实现ViT最经典的变体 ViT-B/16（Base版，16×16的块大小），代码清晰易懂，还加了注释。

3.1 第一步：把图像切成“视觉词汇”

import torch
import torch.nn as nn
import torch.nn.functional as F

class PatchEmbedding(nn.Module):
    """
    图像分块嵌入：用卷积高效实现“分块+线性投影”
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2  # 14×14=196个块
        
        # 卷积核大小=stride=patch_size，一步到位分块+投影
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # 输入: (B, C, H, W) → 输出: (B, n_patches, embed_dim)
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

# 测试一下
if __name__ == "__main__":
    patch_embed = PatchEmbedding()
    dummy_img = torch.randn(2, 3, 224, 224)  # 2张RGB图
    print(f"Patch嵌入后形状: {patch_embed(dummy_img).shape}")  # 输出: torch.Size([2, 196, 768])

3.2 第二步：搭Transformer编码器块

class MultiHeadAttention(nn.Module):
    """
    多头自注意力：简化版实现
    """
    def __init__(self, embed_dim=768, n_heads=12, dropout=0.1):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = embed_dim // n_heads
        self.scale = self.head_dim ** -0.5
        
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.proj = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, N, C = x.shape
        # 1. 计算QKV并拆分多头
        qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)
        
        # 2. 缩放点积注意力
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = self.dropout(attn.softmax(dim=-1))
        
        # 3. 拼接多头并投影
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.proj(x)

class TransformerBlock(nn.Module):
    """
    Transformer编码器块：Pre-Norm结构
    """
    def __init__(self, embed_dim=768, n_heads=12, mlp_ratio=4, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadAttention(embed_dim, n_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        
        # MLP：中间层维度是embed_dim的4倍
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(int(embed_dim * mlp_ratio), embed_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))  # 残差连接1
        return x + self.mlp(self.norm2(x))  # 残差连接2

3.3 第三步：组装完整ViT模型

class VisionTransformer(nn.Module):
    """
    完整ViT-B/16模型
    """
    def __init__(self, img_size=224, n_classes=1000, depth=12):
        super().__init__()
        self.patch_embed = PatchEmbedding()
        
        # CLS Token + 可学习位置编码
        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
        self.pos_embed = nn.Parameter(torch.zeros(1, 196 + 1, 768))  # +1是CLS Token
        self.pos_drop = nn.Dropout(0.1)
        
        # 堆叠12层Transformer块
        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(depth)])
        
        # 分类头
        self.norm = nn.LayerNorm(768)
        self.head = nn.Linear(768, n_classes)

    def forward(self, x):
        B = x.shape[0]
        
        # 1. 图像转Patch
        x = self.patch_embed(x)
        
        # 2. 拼接CLS Token
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        # 3. 加位置编码
        x = self.pos_drop(x + self.pos_embed)
        
        # 4. 通过Transformer编码器
        for block in self.blocks:
            x = block(x)
        
        # 5. 取CLS Token输出分类
        return self.head(self.norm(x)[:, 0])

# 测试完整模型
if __name__ == "__main__":
    vit = VisionTransformer()
    dummy_img = torch.randn(4, 3, 224, 224)
    print(f"ViT输出形状: {vit(dummy_img).shape}")  # 输出: torch.Size([4, 1000])

4. 避坑指南：用好ViT的3个关键

ViT没有CNN的“局部归纳偏置”，**中小规模数据集（<100万张）上直接训ViT必不如CNN**。 ✅ 正确做法：用大规模预训练权重（如ImageNet-21k、MAE）微调！

4.1 超参数选择

超参数推荐配置说明
优化器AdamW (lr=1e-3, weight_decay=0.05)必须用带权重衰减的AdamW
学习率调度Warmup (10 epochs) + Cosine Decay训练初期用小lr，避免CLS Token震荡
数据增强RandAugment + CutMix + MixUp高级增强是ViT在小数据集上收敛的关键
Batch Size越大越好（至少256，推荐1024+）大Batch能稳定注意力权重的训练

4.2 什么时候用ViT，什么时候用CNN？

def choose_between_vit_cnn(data_size, is_speed_critical):
    if data_size < 100_000:
        return "首选CNN（ResNet/EfficientNet），可考虑用DeiT蒸馏"
    elif data_size > 1_000_000 and not is_speed_critical:
        return "首选ViT（或Swin Transformer），用大规模预训练权重微调"
    else:
        return "折中方案：MobileViT（移动端）/CoAtNet（混合架构）"

5. 总结

Vision Transformer用“序列建模”的统一范式，为计算机视觉打开了新的大门。尽管它有“数据 hungry”、“计算量大”的缺点，但在大规模预训练+下游微调的模式下，已经成为图像分类、目标检测、分割等任务的主流选择。

如果想进一步学习ViT的变体，推荐按这个顺序：

  1. DeiT：解决ViT在小数据集上的训练问题
  2. Swin Transformer：引入层次化结构，适配检测/分割任务
  3. MAE：自监督预训练的代表作，大幅降低ViT的数据需求

相关教程

1. 用`timm`库加载预训练的ViT-B/16，在CIFAR-10上微调看看效果 2. 尝试可视化ViT的注意力热力图，直观理解它在关注什么

🔗 扩展阅读