Transformer架构及其在计算机视觉中的革命：从ViT到Swin Transformer
目录
2. 技术基础：自注意力机制与Vision Transformer
2.2 Vision Transformer (ViT)：开创性的直接应用
3.1 Data-efficient Image Transformer (DeiT)：告别大规模预训练
3.2 Swin Transformer：引领视觉任务的通用Backbone
1. Vision Transformer (ViT) 完整实现
2. DeiT (Data-efficient Image Transformer) 实现
摘要：本文深入探讨了Transformer架构如何颠覆性地改变计算机视觉（CV）领域的发展范式。始于自然语言处理（NLP）的Transformer模型，通过Vision Transformer (ViT)、Data-efficient Image Transformer (DeiT) 和 Swin Transformer等一系列里程碑式的工作，成功解决了卷积神经网络（CNN）固有的归纳偏置限制，实现了全局建模能力的突破。本文系统分析了这些核心模型的技术创新、理论基础及其对CV各个子领域的深远影响，并展望了未来多模态统一架构的发展趋势。
关键词：Transformer；计算机视觉；自注意力机制；Vision Transformer；Swin Transformer；全局建模
1. 引言：一场跨领域的架构迁移
2017年，Vaswani等人提出的Transformer架构【1】凭借其强大的自注意力机制和并行化训练能力，彻底改变了自然语言处理领域。然而，计算机视觉长期以来一直是卷积神经网络（CNN） 的天下，其固有的局部连接、权重共享和平移不变性等归纳偏置（Inductive Bias）被证明极其有效。
但CNN也存在本质局限：感受野有限，难以建立真正的长距离依赖关系；层次化特征提取过程可能丢失细节信息。Transformer的核心优势——全局建模能力和动态权重分配——恰好针对这些痛点。本文将详细论述Transformer如何跨越模态鸿沟，引领计算机视觉的新革命。
2. 技术基础：自注意力机制与Vision Transformer
2.1 自注意力机制核心思想
自注意力机制允许序列中的每个元素（"词"或"图像块"）与所有其他元素直接交互，计算其相关性。其核心计算公式为：
Attention(Q, K, V) = softmax(QKᵀ/√dₖ)V
其中Q（Query）、K（Key）、V（Value）均由输入线性变换得到。这种机制实现了：
-
全局感受野：每个位置都能直接关注所有其他位置
-
动态特征权重：注意力权重根据输入内容动态计算
-
强可解释性：注意力图谱可视化模型关注区域
2.2 Vision Transformer (ViT)：开创性的直接应用
2020年，Dosovitskiy等人提出的ViT【2】是首个完全摒弃卷积、纯Transformer架构的CV模型。
-
核心创新：
-
图像分块嵌入：将输入图像分割成固定大小的 patches（如16x16），展平后线性投影为序列 tokens，加入位置编码。
-
类别Token：借鉴BERT的[CLS] token，用于最终分类。
-
Transformer编码器：直接使用标准Transformer编码器堆叠。
-
-
关键发现：
-
大规模数据需求：在JFT-300M等超大数据集上预训练后，ViT性能超越当时最佳CNN。
-
** scalability**：模型越大，性能越好，几乎没有饱和迹象。
-
全局建模优势：在处理需要全局信息的任务（如场景分类、细粒度识别）上表现出色。
-
-
局限性：
-
缺乏CNN固有的平移不变性和局部性归纳偏置。
-
需要大规模预训练数据，在小数据集上容易过拟合。
-
计算复杂度与图像尺寸成平方关系（O(n²)），处理高分辨率图像开销大。
-
ViT证明了"Everything is a sequence" 的范式在视觉领域同样可行，是CV领域的范式转换点。
3. 关键演进：解决ViT的痛点
3.1 Data-efficient Image Transformer (DeiT)：告别大规模预训练
Touvron等人提出的DeiT【3】解决了ViT对海量数据的依赖问题。
-
核心创新：
-
知识蒸馏：引入一个蒸馏token，与class token一起通过注意力机制与教师CNN（如RegNet）交互学习。提出硬标签蒸馏和软标签蒸馏策略。
-
数据增强：结合了强数据增强策略（如Rand-Augment, Mixup, CutMix）。
-
优化策略调整：改进了优化器和超参数设置。
-
-
革命性意义：
-
仅使用ImageNet-1k数据即可训练出高性能ViT，使其应用门槛大幅降低。
-
提供了Transformer+CNN协同的新思路，教师CNN的归纳偏置通过蒸馏传递给学生Transformer。
-
3.2 Swin Transformer：引领视觉任务的通用Backbone
Liu等人提出的Swin Transformer【4】是CV Transformer发展的又一个里程碑，成为了诸多视觉任务的通用骨干网络。
-
核心创新：
-
层级化结构：像CNN一样构建层次化特征图（分辨率逐渐降低，通道数增加），而非ViT的单尺度序列。使其能与现有CV任务 pipeline（如FPN、U-Net）无缝衔接。
-
滑动窗口注意力：在不重叠的局部窗口内计算自注意力，将计算复杂度从O(n²) 降至O(n)。
-
移位窗口：下一层窗口相对上一层进行移位，使不同窗口间能够交互，实现了跨窗口连接，在保持效率的同时获得了全局建模能力。
-
-
革命性意义：
-
效率与性能的完美平衡：解决了ViT计算复杂度高的问题，并能处理高分辨率图像。
-
成为通用Backbone：在分类、检测、分割等任务上全面达到SOTA，证明了其作为CV基础模型的潜力。
-
统一架构的曙光：其设计思想影响了后续诸多工作，推动了视觉乃至多模态通用架构的发展。
-
4. 深远影响与领域革命
4.1 对计算机视觉子领域的颠覆
-
图像分类：ViT/DeiT/Swin系列模型在ImageNet上占据统治地位。
-
目标检测：Swin Transformer作为Backbone的检测器（如Mask R-CNN + Swin）在COCO数据集上取得显著提升。
-
语义分割：基于Swin的U-PerNet等模型在ADE20K等数据集上刷新纪录，得益于其多尺度特征表示。
-
图像生成：Transformer被用于自回归图像生成（如iGPT）和扩散模型（如DiT），展现出强大的序列建模能力。
-
视频理解：时空Transformer（TimeSformer等）通过时间、空间两个维度的注意力，有效建模视频中的长时序依赖。
4.2 研究范式的转变
-
从归纳偏置到规模化数据驱动：减弱了对人工设计偏置（如卷积的局部性）的依赖，更相信模型能从大规模数据中自行学习所需偏置。
-
架构统一趋势：NLP和CV的底层架构趋于统一（均为Transformer），为多模态学习（如CLIP, DALL-E）奠定了坚实基础，推动了“一个模型处理一切”的愿景。
-
Scaling Law的验证：CV领域也验证了模型规模、数据规模和计算量同步增长的“缩放定律”，进入了大模型时代。
5. 挑战与未来展望
尽管革命性巨大，视觉Transformer仍面临挑战：
-
计算复杂度：虽然Swin等模型已优化，但全局注意力仍是计算密集型操作。
-
解释性：尽管有注意力图，但其决策过程相比CNN仍更复杂难解。
-
落地部署：模型参数量大，对硬件要求高，轻量化部署是难题。
未来展望：
-
更高效的注意力机制：如线性注意力、稀疏注意力等将继续发展。
-
神经架构搜索：自动搜索更适合视觉任务的Transformer变体。
-
与CNN的进一步融合：探索hybrid模型，结合二者优势。
-
通往AGI的多模态基础模型：Transformer已成为构建视觉-语言-音频统一大模型的核心基石。
6. 结论
Transformer架构在计算机视觉中的迁移是一场深刻的技术革命。从ViT证明其可行性，到DeiT解决数据依赖，再到Swin Transformer成为通用骨干，这一历程不仅带来了一系列性能卓越的模型，更根本地改变了计算机视觉的研究范式。它打破了视觉与语言的模态壁垒，为构建统一、强大的多模态智能系统铺平了道路，标志着我们向通用人工智能（AGI）又迈进了坚实的一步。
参考文献
[1] Vaswani, A. et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems.
[2] Dosovitskiy, A. et al. (2020). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. International Conference on Learning Representations (ICLR).
[3] Touvron, H. et al. (2021). Training data-efficient image transformers & distillation through attention. International Conference on Machine Learning (ICML).
[4] Liu, Z. et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. IEEE/CVF International Conference on Computer Vision (ICCV).
[5] Carion, N. et al. (2020). End-to-End Object Detection with Transformers. European Conference on Computer Vision (ECCV).
附录：核心模型对比表
| 特性 | ViT | DeiT | Swin Transformer |
|---|---|---|---|
| 核心贡献 | 证明纯Transformer可用于CV | 无需大数据预训练 | 通用、高效的Backbone |
| 关键创新 | 图像分块、类别Token | 蒸馏Token、强数据增强 | 层级结构、滑动窗口、移位窗口 |
| 数据需求 | 极高（JFT-300M） | 低（ImageNet-1K） | 中等（ImageNet-1K） |
| 计算复杂度 | O(n²) | O(n²) | O(n) |
| 主要优势 | 全局建模、Scalability | 数据效率、训练速度 | 多尺度特征、任务通用性 |
| 主要局限 | 数据依赖、计算开销 | - | 相对复杂的设计 |
| 影响力 | 范式开创者 | 实用化推动者 | 领域统一者 |
1. Vision Transformer (ViT) 完整实现
python
复制
下载
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from einops import rearrange, repeat
class PatchEmbedding(nn.Module):
"""将图像分割成patches并嵌入"""
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.num_patches = (img_size // patch_size) ** 2
self.proj = nn.Conv2d(in_channels, embed_dim,
kernel_size=patch_size, stride=patch_size)
def forward(self, x):
# x: [B, C, H, W]
x = self.proj(x) # [B, E, H/p, W/p]
x = x.flatten(2) # [B, E, N]
x = x.transpose(1, 2) # [B, N, E]
return x
class PositionalEncoding(nn.Module):
"""可学习的位置编码"""
def __init__(self, num_patches, embed_dim):
super().__init__()
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
def forward(self, x):
# x: [B, N, E]
batch_size = x.size(0)
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
x = torch.cat((cls_tokens, x), dim=1)
x = x + self.pos_embed
return x
class MultiHeadSelfAttention(nn.Module):
"""多头自注意力机制"""
def __init__(self, embed_dim=768, num_heads=12, dropout=0.0):
super().__init__()
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
assert self.head_dim * num_heads == embed_dim
self.qkv = nn.Linear(embed_dim, embed_dim * 3)
self.attn_drop = nn.Dropout(dropout)
self.proj = nn.Linear(embed_dim, embed_dim)
self.proj_drop = nn.Dropout(dropout)
def forward(self, x, mask=None):
batch_size, seq_len, embed_dim = x.shape
# 生成Q, K, V
qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4) # [3, B, H, N, D]
q, k, v = qkv[0], qkv[1], qkv[2]
# 计算注意力分数
attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
if mask is not None:
attn = attn.masked_fill(mask == 0, -1e9)
attn = F.softmax(attn, dim=-1)
attn = self.attn_drop(attn)
# 注意力加权
x = (attn @ v).transpose(1, 2).reshape(batch_size, seq_len, embed_dim)
x = self.proj(x)
x = self.proj_drop(x)
return x, attn
class MLP(nn.Module):
"""多层感知机"""
def __init__(self, in_features, hidden_features=None, out_features=None, dropout=0.0):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = nn.GELU()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(dropout)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class TransformerBlock(nn.Module):
"""Transformer编码块"""
def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.norm1 = nn.LayerNorm(embed_dim)
self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
self.norm2 = nn.LayerNorm(embed_dim)
mlp_hidden_dim = int(embed_dim * mlp_ratio)
self.mlp = MLP(embed_dim, mlp_hidden_dim, dropout=dropout)
def forward(self, x, mask=None):
# 自注意力子层
attn_output, attn_weights = self.attn(self.norm1(x), mask)
x = x + attn_output
# MLP子层
x = x + self.mlp(self.norm2(x))
return x, attn_weights
class VisionTransformer(nn.Module):
"""完整的ViT模型"""
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, emb_dropout=0.1):
super().__init__()
self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
num_patches = self.patch_embed.num_patches
self.pos_embed = PositionalEncoding(num_patches, embed_dim)
self.dropout = nn.Dropout(emb_dropout)
# Transformer编码器
self.blocks = nn.ModuleList([
TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout)
for _ in range(depth)
])
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
self.initialize_weights()
def initialize_weights(self):
# 初始化权重
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.LayerNorm):
nn.init.zeros_(m.bias)
nn.init.ones_(m.weight)
# 特殊初始化位置编码
nn.init.trunc_normal_(self.pos_embed.pos_embed, std=0.02)
nn.init.trunc_normal_(self.pos_embed.cls_token, std=0.02)
def forward(self, x, return_attn=False):
# 图像分块嵌入
x = self.patch_embed(x)
# 添加位置编码和cls token
x = self.pos_embed(x)
x = self.dropout(x)
# 通过Transformer块
attn_weights = []
for blk in self.blocks:
x, attn = blk(x)
if return_attn:
attn_weights.append(attn)
# 分类头
x = self.norm(x)
cls_token = x[:, 0] # 取cls token对应的特征
logits = self.head(cls_token)
if return_attn:
return logits, attn_weights
return logits
# ViT模型配置
def vit_tiny(patch_size=16, **kwargs):
return VisionTransformer(patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, **kwargs)
def vit_small(patch_size=16, **kwargs):
return VisionTransformer(patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, **kwargs)
def vit_base(patch_size=16, **kwargs):
return VisionTransformer(patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, **kwargs)
def vit_large(patch_size=16, **kwargs):
return VisionTransformer(patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, **kwargs)
2. DeiT (Data-efficient Image Transformer) 实现
python
复制
下载
class DistillationToken(nn.Module):
"""DeiT蒸馏token"""
def __init__(self, embed_dim=768):
super().__init__()
self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
def forward(self, x):
batch_size = x.size(0)
dist_tokens = self.dist_token.expand(batch_size, -1, -1)
x = torch.cat((x, dist_tokens), dim=1)
return x
class DeiT(VisionTransformer):
"""DeiT模型，支持知识蒸馏"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# 添加蒸馏token
self.dist_token = nn.Parameter(torch.zeros(1, 1, self.patch_embed.embed_dim))
num_patches = self.patch_embed.num_patches
self.pos_embed = nn.Parameter(
torch.zeros(1, num_patches + 2, self.patch_embed.embed_dim))
# 蒸馏头
self.head_dist = nn.Linear(self.patch_embed.embed_dim, self.num_classes)
# 重新初始化
nn.init.trunc_normal_(self.dist_token, std=0.02)
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.zeros_(self.head_dist.weight)
def forward(self, x):
# 图像分块嵌入
B = x.shape[0]
x = self.patch_embed(x)
# 添加cls token和dist token
cls_tokens = self.pos_embed.cls_token.expand(B, -1, -1)
dist_tokens = self.dist_token.expand(B, -1, -1)
x = torch.cat((cls_tokens, x, dist_tokens), dim=1)
# 添加位置编码
x = x + self.pos_embed
x = self.dropout(x)
# Transformer编码
for blk in self.blocks:
x, _ = blk(x)
x = self.norm(x)
# 两个输出：cls token和dist token
x_cls = self.head(x[:, 0])
x_dist = self.head_dist(x[:, -1])
if self.training:
return x_cls, x_dist
else:
# 推理时平均两个输出
return (x_cls + x_dist) / 2
class DeiTWithTeacher(nn.Module):
"""完整的DeiT训练框架，包含教师模型"""
def __init__(self, student_model, teacher_model):
super().__init__()
self.student = student_model
self.teacher = teacher_model
self.teacher.eval() # 教师模型固定
def forward(self, x, labels=None):
student_logits, student_dist_logits = self.student(x)
with torch.no_grad():
teacher_logits = self.teacher(x)
if labels is not None:
# 计算损失：学生vs真实标签 + 学生vs教师
loss_cls = F.cross_entropy(student_logits, labels)
loss_dist = F.kl_div(
F.log_softmax(student_dist_logits / 2.0, dim=-1),
F.softmax(teacher_logits / 2.0, dim=-1),
reduction='batchmean'
) * (2.0 * 2.0)
# 硬标签蒸馏
teacher_pred = torch.argmax(teacher_logits, dim=1)
loss_hard = F.cross_entropy(student_dist_logits, teacher_pred)
total_loss = loss_cls + 0.5 * loss_dist + 0.5 * loss_hard
return total_loss, student_logits
return student_logits
3. Swin Transformer 核心组件实现
python
复制
下载
class WindowAttention(nn.Module):
"""基于窗口的多头自注意力"""
def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.):
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))
# 生成相对位置索引
coords_h = torch.arange(self.window_size[0])
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w]))
coords_flatten = torch.flatten(coords, 1)
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
relative_coords = relative_coords.permute(1, 2, 0).contiguous()
relative_coords[:, :, 0] += self.window_size[0] - 1
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1)
self.register_buffer("relative_position_index", relative_position_index)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
nn.init.trunc_normal_(self.relative_position_bias_table, std=.02)
def forward(self, x, mask=None):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) * self.scale
# 添加相对位置偏置
relative_position_bias = self.relative_position_bias_table[
self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1],
self.window_size[0] * self.window_size[1], -1)
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
attn = attn + relative_position_bias.unsqueeze(0)
if mask is not None:
nW = mask.shape[0]
attn = attn.view(B // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, N, N)
attn = F.softmax(attn, dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
def window_partition(x, window_size):
"""将特征图划分成窗口"""
B, H, W, C = x.shape
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
return windows
def window_reverse(windows, window_size, H, W):
"""将窗口重组为特征图"""
B = int(windows.shape[0] / (H * W / window_size / window_size))
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
return x
class SwinTransformerBlock(nn.Module):
"""Swin Transformer块"""
def __init__(self, dim, input_resolution, num_heads, window_size=7,
shift_size=0, mlp_ratio=4., drop=0., attn_drop=0.):
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_size:
self.shift_size = 0
self.window_size = min(self.input_resolution)
self.norm1 = nn.LayerNorm(dim)
self.attn = WindowAttention(
dim, window_size=(self.window_size, self.window_size),
num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop)
self.norm2 = nn.LayerNorm(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = MLP(dim, mlp_hidden_dim, drop=drop)
if self.shift_size > 0:
# 生成注意力掩码
H, W = self.input_resolution
img_mask = torch.zeros((1, H, W, 1))
h_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
cnt = 0
for h in h_slices:
for w in w_slices:
img_mask[:, h, w, :] = cnt
cnt += 1
mask_windows = window_partition(img_mask, self.window_size)
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0))
attn_mask = attn_mask.masked_fill(attn_mask == 0, float(0.0))
else:
attn_mask = None
self.register_buffer("attn_mask", attn_mask)
def forward(self, x):
H, W = self.input_resolution
B, L, C = x.shape
assert L == H * W, "input feature has wrong size"
shortcut = x
x = self.norm1(x)
x = x.view(B, H, W, C)
# 循环移位
if self.shift_size > 0:
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
else:
shifted_x = x
# 窗口划分
x_windows = window_partition(shifted_x, self.window_size)
x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
# 窗口注意力
attn_windows = self.attn(x_windows, mask=self.attn_mask)
# 窗口重组
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
shifted_x = window_reverse(attn_windows, self.window_size, H, W)
# 反向循环移位
if self.shift_size > 0:
x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
x = shifted_x
x = x.view(B, H * W, C)
x = shortcut + x
# MLP
x = x + self.mlp(self.norm2(x))
return x
4. 训练示例代码
python
复制
下载
def train_vit_model():
"""ViT训练示例"""
# 初始化模型
model = vit_base(num_classes=1000, img_size=224)
# 数据增强
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 优化器设置
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.05)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
criterion = nn.CrossEntropyLoss()
# 训练循环
model.train()
for epoch in range(300):
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f'Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item():.4f}')
scheduler.step()
# 模型测试
def test_model_performance(model, test_loader):
"""测试模型性能"""
model.eval()
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
_, predicted = torch.max(output.data, 1)
total += target.size(0)
correct += (predicted == target).sum().item()
accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')
return accuracy
# 注意力可视化
def visualize_attention(model, image):
"""可视化注意力图"""
model.eval()
with torch.no_grad():
_, attn_weights = model(image.unsqueeze(0), return_attn=True)
# 选择最后一层的注意力图
last_layer_attn = attn_weights[-1]
# 平均所有头的注意力
attn_map = last_layer_attn.mean(dim=1)[0, 0, 1:] # 取cls token对其他patch的注意力
# 重塑为2D注意力图
patch_size = model.patch_embed.patch_size
h = w = int(math.sqrt(attn_map.shape[0]))
attn_map = attn_map.reshape(h, w)
return attn_map
if __name__ == "__main__":
# 示例用法
model = vit_base(num_classes=1000)
print(f"ViT参数量: {sum(p.numel() for p in model.parameters()):,}")
# 测试前向传播
dummy_input = torch.randn(1, 3, 224, 224)
output = model(dummy_input)
print(f"输入形状: {dummy_input.shape}")
print(f"输出形状: {output.shape}")
这些代码提供了完整的Transformer视觉模型实现，包含了：
-
ViT：基础Vision Transformer架构
-
DeiT：包含知识蒸馏的改进版本
-
Swin Transformer：层级化和窗口注意力机制
-
训练框架：完整的训练和评估逻辑
更多推荐
16
0- 0
已为社区贡献7条内容
所有评论(0)