PyTorch深度学习实战 |手算ViT(Vision Transformer)模型

# ============================== 1. 环境依赖导入 ==============================
import torch
import torch.nn as nn
import torch.nn.functional as F
# ============================== 2. ViT核心组件定义（修复版） ==============================
class PatchEmbedding(nn.Module):
    """图像块嵌入：将图像分割为块并映射为Token序列"""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2  # 14×14=196（Patch Token数量）
        self.conv = nn.Conv2d(
            in_channels=in_channels,
            out_channels=embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
    def forward(self, x):
        # 输入x：[batch_size, 3, 224, 224]（PyTorch图像格式）
        print(f"PatchEmbedding输入形状: {x.shape}")
        # 卷积：[B, 3, 224, 224] → [B, 768, 14, 14]（每个像素对应1个16×16图像块）
        x = self.conv(x)
        print(f"PatchEmbedding卷积后形状: {x.shape}")
        # 展平：[B, 768, 14, 14] → [B, 768, 196]（196=14×14，按列展平）
        x = x.flatten(2)
        print(f"PatchEmbedding展平后形状: {x.shape}")
        # 转置：[B, 768, 196] → [B, 196, 768]（适配Transformer的[B, N, D]输入格式）
        x = x.transpose(1, 2)
        print(f"PatchEmbedding输出形状: {x.shape}\n")
        return x
class MultiHeadSelfAttention(nn.Module):
    """多头自注意力：捕捉Token间全局依赖关系"""
    def __init__(self, embed_dim=768, num_heads=12, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads  # 768÷12=64（单个头的维度）
        assert self.head_dim * num_heads == embed_dim, "embed_dim需被num_heads整除"
        
        self.qkv = nn.Linear(embed_dim, embed_dim * 3)  # 一次性生成Q/K/V
        self.proj = nn.Linear(embed_dim, embed_dim)      # 多头结果融合
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        # 输入x：[B, 197, 768]（196个Patch Token + 1个CLS Token）
        batch_size, num_tokens, embed_dim = x.shape
        print(f"MHSA输入形状: {x.shape}")
        
        # 1. 生成Q/K/V：[B, 197, 768] → [B, 197, 2304] → 拆分为3个[B, 197, 768]
        qkv = self.qkv(x).reshape(batch_size, num_tokens, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # 每个形状：[B, 12, 197, 64]
        print(f"MHSA中Q/K/V形状（单头）: {q.shape}")
        
        # 2. 计算注意力权重：Q×K^T / √d_k → [B, 12, 197, 197]（每行是1个Token对所有Token的注意力）
        attn_weights = torch.matmul(q, k.transpose(-2, -1))
        attn_weights = attn_weights / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        print(f"MHSA注意力权重形状: {attn_weights.shape}")
        
        # 3. Softmax归一化 + Dropout：确保权重和为1，防止过拟合
        attn_weights = F.softmax(attn_weights, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # 4. 注意力加权V：[B, 12, 197, 197] × [B, 12, 197, 64] → [B, 12, 197, 64]
        attn_output = torch.matmul(attn_weights, v)
        print(f"MHSA单头输出形状: {attn_output.shape}")
        
        # 5. 多头拼接：[B, 12, 197, 64] → [B, 197, 768]（12×64=768，恢复原维度）
        attn_output = attn_output.permute(0, 2, 1, 3).reshape(batch_size, num_tokens, embed_dim)
        print(f"MHSA多头拼接后形状: {attn_output.shape}")
        
        # 6. 融合输出：[B, 197, 768] → [B, 197, 768]（线性层调整特征）
        x = self.proj(attn_output)
        x = self.dropout(x)
        print(f"MHSA最终输出形状: {x.shape}\n")
        return x
class TransformerEncoderLayer(nn.Module):
    """Transformer编码器层：LN→MHSA→残差 + LN→MLP→残差（Pre-LN结构）"""
    def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)  # MHSA前的归一化
        self.norm2 = nn.LayerNorm(embed_dim)  # MLP前的归一化
        self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, int(embed_dim * mlp_ratio)),  # 768→3072（维度扩张）
            nn.GELU(),                                          # 激活函数（比ReLU更平滑）
            nn.Dropout(dropout),
            nn.Linear(int(embed_dim * mlp_ratio), embed_dim),  # 3072→768（维度收缩）
            nn.Dropout(dropout)
        )
    def forward(self, x):
        # 输入x：[B, 197, 768]（来自上一层编码器或输入层）
        print(f"EncoderLayer输入形状: {x.shape}")
        
        # 1. MHSA分支：残差连接（x + MHSA(LN(x))）→ 保持维度不变
        x = x + self.attn(self.norm1(x))
        print(f"EncoderLayer MHSA+残差后形状: {x.shape}")
        
        # 2. MLP分支：残差连接（x + MLP(LN(x))）→ 保持维度不变
        x = x + self.mlp(self.norm2(x))
        print(f"EncoderLayer MLP+残差后形状: {x.shape}\n")
        return x
class VisionTransformer(nn.Module):
    """完整ViT模型：PatchEmbedding→位置嵌入+CLS Token→编码器堆叠→分类头（修复pos_embed维度）"""
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768,
                 num_heads=12, num_layers=12, num_classes=1000, dropout=0.1):
        super().__init__()
        # 1. 图像块嵌入模块
        self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
        num_patches = self.patch_embed.num_patches  # 196（Patch Token数量）
        
        # 2. 修复：位置嵌入长度改为197（适配CLS Token + Patch Token）
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))  # [1, 197, 768]
        # 3. CLS Token（用于聚合全局特征）
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))  # [1, 1, 768]
        self.dropout = nn.Dropout(dropout)
        
        # 4. 12层Transformer编码器堆叠
        self.encoders = nn.ModuleList([
            TransformerEncoderLayer(embed_dim, num_heads, mlp_ratio=4.0, dropout=dropout)
            for _ in range(num_layers)
        ])
        
        # 5. 输出层分类头（LN+线性层）
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, num_classes)
    def forward(self, x):
        batch_size = x.shape[0]
        # 输入x：[B, 3, 224, 224]（原始RGB图像）
        print(f"ViT整体输入形状: {x.shape}\n")
        
        # -------------------------- 步骤1：图像→Patch Token序列 --------------------------
        x = self.patch_embed(x)  # 输出：[B, 196, 768]
        
        # -------------------------- 步骤2：插入CLS Token --------------------------
        # 扩展CLS Token到batch维度：[1, 1, 768] → [B, 1, 768]
        cls_token = self.cls_token.expand(batch_size, -1, -1)
        # 拼接：[B, 1, 768] + [B, 196, 768] → [B, 197, 768]
        x = torch.cat((cls_token, x), dim=1)
        print(f"插入CLS Token后形状: {x.shape}\n")
        
        # -------------------------- 步骤3：添加位置嵌入（修复后可正常相加） --------------------------
        # 现在x和pos_embed形状均为[B, 197, 768]，可直接广播相加
        x = x + self.pos_embed
        x = self.dropout(x)
        print(f"添加位置嵌入后形状: {x.shape}")
        print(f"Dropout后形状: {x.shape}\n")
        
        # -------------------------- 步骤4：12层编码器特征提取 --------------------------
        for i, encoder in enumerate(self.encoders):
            print(f"=== 第{i+1}层编码器 ===")
            x = encoder(x)  # 每层输入输出均为[B, 197, 768]
        
        # -------------------------- 步骤5：分类预测 --------------------------
        # 提取CLS Token：[B, 197, 768] → [B, 768]（仅用第0个Token的全局特征）
        cls_token_final = x[:, 0, :]
        print(f"\n提取CLS Token后形状: {cls_token_final.shape}")
        
        # LN归一化：[B, 768] → [B, 768]（稳定分类头输入）
        cls_token_final = self.norm(cls_token_final)
        print(f"LN归一化后形状: {cls_token_final.shape}")
        
        # 分类头：[B, 768] → [B, 1000]（1000为ImageNet-1K类别数）
        logits = self.head(cls_token_final)
        print(f"分类头输出形状: {logits.shape}")
        
        return logits
# ============================== 3. 模型测试（可直接运行，无维度报错） ==============================
if __name__ == "__main__":
    # 1. 设置随机种子（确保结果可复现）
    torch.manual_seed(42)
    
    # 2. 构造测试输入（batch_size=1，RGB图像：[B, C, H, W]）
    dummy_img = torch.randn(1, 3, 224, 224)  # 随机生成符合格式的图像
    
    # 3. 实例化ViT-Base模型（参数与论文一致）
    vit_model = VisionTransformer(
        img_size=224,
        patch_size=16,
        in_channels=3,
        embed_dim=768,
        num_heads=12,
        num_layers=12,
        num_classes=1000
    )
    print(f"ViT-Base模型初始化完成（参数总数：{sum(p.numel() for p in vit_model.parameters()):,}）\n")
    
    # 4. 前向传播测试（禁用梯度计算，仅推理）
    with torch.no_grad():
        output = vit_model(dummy_img)
    
    # 5. 输出结果分析
    probabilities = F.softmax(output, dim=1)  # 转化为类别概率（和为1）
    pred_class = torch.argmax(probabilities, dim=1).item()  # 预测概率最高的类别索引
    
    print(f"\n=== 最终结果 ===")
    print(f"模型输出logits形状: {output.shape}")          # [1, 1000]
    print(f"类别概率分布形状: {probabilities.shape}")    # [1, 1000]
    print(f"预测概率最高的类别索引: {pred_class}")       # 随机值（模型未训练）