Transformer架构及其在计算机视觉中的革命：从ViT到Swin Transformer

目录

1. 引言：一场跨领域的架构迁移

2. 技术基础：自注意力机制与Vision Transformer

2.1 自注意力机制核心思想

2.2 Vision Transformer (ViT)：开创性的直接应用

3. 关键演进：解决ViT的痛点

3.1 Data-efficient Image Transformer (DeiT)：告别大规模预训练

3.2 Swin Transformer：引领视觉任务的通用Backbone

4. 深远影响与领域革命

4.1 对计算机视觉子领域的颠覆

4.2 研究范式的转变

5. 挑战与未来展望

6. 结论

参考文献

1. Vision Transformer (ViT) 完整实现

2. DeiT (Data-efficient Image Transformer) 实现

3. Swin Transformer 核心组件实现

4. 训练示例代码

---

摘要：本文深入探讨了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）均由输入线性变换得到。这种机制实现了：

1. 全局感受野：每个位置都能直接关注所有其他位置

2. 动态特征权重：注意力权重根据输入内容动态计算

3. 强可解释性：注意力图谱可视化模型关注区域

2.2 Vision Transformer (ViT)：开创性的直接应用

2020年，Dosovitskiy等人提出的ViT【2】是首个完全摒弃卷积、纯Transformer架构的CV模型。

• 核心创新：

  1. 图像分块嵌入：将输入图像分割成固定大小的 patches（如16x16），展平后线性投影为序列 tokens，加入位置编码。

  2. 类别Token：借鉴BERT的[CLS] token，用于最终分类。

  3. 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对海量数据的依赖问题。

• 核心创新：

  1. 知识蒸馏：引入一个蒸馏token，与class token一起通过注意力机制与教师CNN（如RegNet）交互学习。提出硬标签蒸馏和软标签蒸馏策略。

  2. 数据增强：结合了强数据增强策略（如Rand-Augment, Mixup, CutMix）。

  3. 优化策略调整：改进了优化器和超参数设置。

• 革命性意义：

  • 仅使用ImageNet-1k数据即可训练出高性能ViT，使其应用门槛大幅降低。

  • 提供了Transformer+CNN协同的新思路，教师CNN的归纳偏置通过蒸馏传递给学生Transformer。

3.2 Swin Transformer：引领视觉任务的通用Backbone

Liu等人提出的Swin Transformer【4】是CV Transformer发展的又一个里程碑，成为了诸多视觉任务的通用骨干网络。

• 核心创新：

  1. 层级化结构：像CNN一样构建层次化特征图（分辨率逐渐降低，通道数增加），而非ViT的单尺度序列。使其能与现有CV任务 pipeline（如FPN、U-Net）无缝衔接。

  2. 滑动窗口注意力：在不重叠的局部窗口内计算自注意力，将计算复杂度从O(n²) 降至O(n)。

  3. 移位窗口：下一层窗口相对上一层进行移位，使不同窗口间能够交互，实现了跨窗口连接，在保持效率的同时获得了全局建模能力。

• 革命性意义：

  • 效率与性能的完美平衡：解决了ViT计算复杂度高的问题，并能处理高分辨率图像。

  • 成为通用Backbone：在分类、检测、分割等任务上全面达到SOTA，证明了其作为CV基础模型的潜力。

  • 统一架构的曙光：其设计思想影响了后续诸多工作，推动了视觉乃至多模态通用架构的发展。

4. 深远影响与领域革命

4.1 对计算机视觉子领域的颠覆

1. 图像分类：ViT/DeiT/Swin系列模型在ImageNet上占据统治地位。

2. 目标检测：Swin Transformer作为Backbone的检测器（如Mask R-CNN + Swin）在COCO数据集上取得显著提升。

3. 语义分割：基于Swin的U-PerNet等模型在ADE20K等数据集上刷新纪录，得益于其多尺度特征表示。

4. 图像生成：Transformer被用于自回归图像生成（如iGPT）和扩散模型（如DiT），展现出强大的序列建模能力。

5. 视频理解：时空Transformer（TimeSformer等）通过时间、空间两个维度的注意力，有效建模视频中的长时序依赖。

4.2 研究范式的转变

1. 从归纳偏置到规模化数据驱动：减弱了对人工设计偏置（如卷积的局部性）的依赖，更相信模型能从大规模数据中自行学习所需偏置。

2. 架构统一趋势：NLP和CV的底层架构趋于统一（均为Transformer），为多模态学习（如CLIP, DALL-E）奠定了坚实基础，推动了“一个模型处理一切”的愿景。

3. Scaling Law的验证：CV领域也验证了模型规模、数据规模和计算量同步增长的“缩放定律”，进入了大模型时代。

5. 挑战与未来展望

尽管革命性巨大，视觉Transformer仍面临挑战：

1. 计算复杂度：虽然Swin等模型已优化，但全局注意力仍是计算密集型操作。

2. 解释性：尽管有注意力图，但其决策过程相比CNN仍更复杂难解。

3. 落地部署：模型参数量大，对硬件要求高，轻量化部署是难题。

未来展望：

1. 更高效的注意力机制：如线性注意力、稀疏注意力等将继续发展。

2. 神经架构搜索：自动搜索更适合视觉任务的Transformer变体。

3. 与CNN的进一步融合：探索hybrid模型，结合二者优势。

4. 通往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视觉模型实现，包含了：

1. ViT：基础Vision Transformer架构

2. DeiT：包含知识蒸馏的改进版本

3. Swin Transformer：层级化和窗口注意力机制

4. 训练框架：完整的训练和评估逻辑