把CNN换到ViT时踩过的坑
把CNN换到ViT时踩过的坑上手并不难，难的是稳定跑起来。
下面只记真正影响结果的部分。
从CNN到ViT的动机
先用一句话说清楚区别：CNN通过滑动窗口在局部提取特征，一层层堆叠后逐渐扩大感受野；ViT则是直接把图像切成一系列图块（patch），然后用Transformer的自注意力机制让每个图块都能和其他所有图块交互。
这个差异在代码层面就很明显。CNN的卷积操作是空间局部化的：
import torch
import torch.nn as nn
# CNN典型写法：局部卷积 + 池化
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.fc = nn.Linear(128 * 8 * 8, 10) # 假设输入32x32
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(x.size(0), -1)
return self.fc(x)
而ViT的逻辑则是先把图切成块，然后直接当作序列处理：
import torch
import torch.nn as nn
import math
class ViTPatchEmbedding(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
super().__init__()
self.img_size = img_size
self.patch_size = patch_size
self.n_patches = (img_size // patch_size) ** 2
# 用卷积把patch展平成向量，实际就是一个convolution with kernel=patch_size
self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
# x: [batch_size, channels, height, width]
x = self.proj(x) # [batch_size, embed_dim, n_patches**0.5, n_patches**0.5]
x = x.flatten(2) # [batch_size, embed_dim, n_patches]
x = x.transpose(1, 2) # [batch_size, n_patches, embed_dim]
return x
这里有个细节很关键：ViT的patch embedding本质上就是一个卷积核大小等于patch_size、stride等于patch_size的卷积操作。所以从实现上看，ViT并不是完全抛弃卷积，而是用了一个非常粗糙的卷积来做初始特征提取。
自注意力机制在图像上的实际效果
自注意力机制是Transformer的核心，在NLP里它让每个词都能关注到序列中的其他词。在视觉任务中，这个机制让每个图块都能和其他所有图块"对话"。
简单实现一个多头自注意力：
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim=768, num_heads=12, dropout=0.1):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_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, mask=None):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) * self.scale
if mask is not None:
attn = attn.masked_fill(mask == 0, float('-inf'))
attn = attn.softmax(dim=-1)
attn = self.dropout(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
return x
在实践过程中，我发现自注意力机制的计算复杂度是O(N²)，其中N是序列长度（也就是图像的图块数量）。对于一个224×224的图像，如果patch_size=16，那么就有196个图块，注意力矩阵大小是196×196，这还能接受。但如果图像更大或者patch更小，计算量就会爆炸。
这也是为什么在早期实践中，ViT在小数据集上表现不如CNN的一个原因：数据量不够时，自注意力的全局建模能力体现不出来，但计算成本却实实在在地增加了。
完整的ViT实现和实际部署
下面是一个简化版的ViT实现，省略了一些细节但保留了核心结构：
class VisionTransformer(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000,
embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, dropout=0.1):
super().__init__()
self.patch_embed = ViTPatchEmbedding(img_size, patch_size, in_channels, embed_dim)
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.patch_embed.n_patches, embed_dim))
self.pos_drop = nn.Dropout(dropout)
self.blocks = nn.ModuleList([
nn.ModuleDict({
'norm1': nn.LayerNorm(embed_dim),
'attn': MultiHeadAttention(embed_dim, num_heads, dropout),
'norm2': nn.LayerNorm(embed_dim),
'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)
)
})
for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
def forward(self, x):
x = self.patch_embed(x)
cls_token = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat((cls_token, x), dim=1)
x = x + self.pos_embed
x = self.pos_drop(x)
for block in self.blocks:
x = x + block['attn'](block['norm1'](x))
x = x + block['mlp'](block['norm2'](x))
x = self.norm(x)
return self.head(x[:, 0]) # 使用cls token的分类结果
在实际部署中，我遇到过几个坑：
第一个是位置编码的初始化。论文里提到使用正弦位置编码或可学习的位置编码，但在实践中我发现可学习位置编码收敛更快，尤其是在数据量不大的情况下。不过要注意，如果推理时的图像尺寸和训练时不同，位置编码需要插值适配：
def resize_pos_embed(posemb, new_grid_size):
cls_token_emb = posemb[:, :1]
grid_emb = posemb[:, 1:].reshape(1, int(math.sqrt(posemb.shape[1] - 1)),
int(math.sqrt(posemb.shape[1] - 1)), -1)
grid_emb = grid_emb.permute(0, 3, 1, 2)
grid_emb = torch.nn.functional.interpolate(
grid_emb, size=new_grid_size, mode='bicubic', align_corners=False)
grid_emb = grid_emb.permute(0, 2, 3, 1).reshape(1, -1, grid_emb.shape[1])
posemb = torch.cat([cls_token_emb, grid_emb], dim=1)
return posemb
第二个坑是计算资源的需求。ViT的显存占用比相同参数量的CNN要高，主要原因就是注意力矩阵。在RTX 3090上训练ViT-Base（约86M参数）时，batch size通常只能开到64左右，而ResNet-50可以轻松开到256。
第三个坑是训练稳定性。ViT对学习率和优化器比较敏感，在实践中我发现用AdamW配合cosine learning rate schedule效果比较稳定：
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
CNN和ViT的实际选择
在实际项目中，选择哪种架构往往不是看谁的精度高，而是看资源约束和场景需求。
CNN的优势在于：
- 计算效率高，适合移动端部署
- 在小数据集上表现更稳定
- 有大量的预训练模型和调参经验
- 工具链成熟，TensorRT、ONNX等支持好
ViT的优势在于：
- 全局建模能力强，适合需要长距离依赖的场景
- 端到端学习能力强，在大规模数据上能学到更好的特征表示
- 架构更简洁，不需要手工设计网络层次
- 迁移学习效果好，尤其是从大规模预训练模型开始
在实践中，我发现一个比较实用的做法是：先用CNN快速迭代和验证想法，等方案稳定后再考虑用ViT进行最终的精度提升。这样既能保证开发效率，又能在最后阶段拿到更好的效果。
数据依赖和预训练策略
ViT对数据量的要求比CNN高得多。在ImageNet-1K上从头训练ViT，效果往往不如ResNet。但如果在更大规模的数据集（如ImageNet-21K、JFT-300M）上预训练，然后再微调，ViT就能展现出优势。
这也是为什么在实际应用中，大家通常都是用预训练的ViT模型进行微调：
from transformers import ViTForImageClassification, ViTImageProcessor
# 加载预训练模型
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
# 微调最后一层分类头
model.classifier = nn.Linear(model.config.hidden_size, num_classes)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
在微调过程中，我发现冻结前面的Transformer层，只训练分类头，对于小数据集往往更稳定：
for name, param in model.named_parameters():
if 'classifier' not in name:
param.requires_grad = False
当前的一些发展和方向
现在视觉Transformer已经发展出很多变体，比如Swin Transformer用层次化结构降低了计算复杂度，DeiT通过知识蒸馏让ViT在ImageNet-1K上也能训练得很好。
Swin Transformer的核心想法是把自注意力限制在局部窗口，然后通过shifted window实现跨窗口信息流动：
class WindowAttention(nn.Module):
def __init__(self, dim, window_size, num_heads):
super().__init__()
self.dim = dim
self.window_size = window_size
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = head_dim ** -0.5
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
self.qkv = nn.Linear(dim, dim * 3)
self.attn_drop = nn.Dropout(0.1)
self.proj = nn.Linear(dim, dim)
def forward(self, x, mask=None):
B_, N, C = x.shape
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
q = q * self.scale
attn = q @ k.transpose(-2, -1)
if mask is not None:
attn = attn + mask.unsqueeze(1)
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
x = self.proj(x)
return x
这种设计显著降低了计算复杂度，同时也保持了建模长距离依赖的能力。
一些实际经验
在折腾ViT的这段时间里，我总结了几条比较实用的经验：
数据量足够（比如ImageNet-1K级别）时，ViT的精度通常会超过CNN。但如果数据量只有几千张图像，CNN可能更稳妥。
计算资源有限时，优先考虑CNN或者轻量级的ViT变体。ViT-Base在1080Ti上训练会比较吃力，而ResNet-50可以轻松跑起来。
如果需要迁移学习，ViT的效果往往比CNN好。特别是在跨域场景下，ViT的泛化能力会更强。
部署时要注意推理优化。ViT的注意力矩阵可以通过Flash Attention等优化技术加速，但需要硬件支持。CNN则可以通过TensorRT等工具直接优化。
调试时可以先用小一点的模型验证思路，比如ViT-Tiny或者ViT-Small，确认方案可行后再上大模型。
结尾
从CNN到ViT，不只是架构的切换，更像是处理图像信息的方式发生了变化。CNN像是一个经验丰富的工匠，手工设计了特征提取的层次；ViT则更像是一个自学的学生，通过大量数据自己学会了如何看图。
这两条路线现在正在融合，ConvNeXt把卷积网络改造成了Transformer的风格，而ViT也在借鉴CNN的层次化设计思路。或许未来的视觉模型会是一个既懂局部特征又懂全局依赖的混合体，但不管怎么变，核心还是理解图像、提取有效信息的能力。
从工程实践的角度看，选择哪种架构最终还是取决于场景、数据和资源约束。没有最好的模型，只有最合适的模型。但多了解一种思路，至少在遇到问题时多了一个选择，这不就是技术人最在乎的吗？
版权声明: 本文首发于 指尖魔法屋-把CNN换到ViT时踩过的坑（https://blog.thinkmoon.cn/post/232-ai-vision-transformer-cnn-to-vit/） 转载或引用必须申明原指尖魔法屋来源及源地址！
评论
使用 GitHub 账号登录后即可留言，支持 Markdown。