vit

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    """
    将输入图像分块（Patch）并进行线性映射
    """
    def __init__(self, in_channels=3, patch_size=16, embed_dim=768, img_size=224):
        super().__init__()
        self.patch_size = patch_size
        self.img_size = img_size
        # 计算图像可以分出多少个 patch
        self.num_patches = (img_size // patch_size) ** 2
        
        # 使用步长为 patch_size 的卷积层来等效实现图像的切块和线性映射
        self.proj = nn.Conv2d(
            in_channels, 
            embed_dim, 
            kernel_size=patch_size, 
            stride=patch_size
        )

    def forward(self, x):
        # x shape: (Batch_size, in_channels, H, W)
        x = self.proj(x)  # shape: (Batch_size, embed_dim, H/patch_size, W/patch_size)
        x = x.flatten(2)  # shape: (Batch_size, embed_dim, num_patches)
        x = x.transpose(1, 2)  # shape: (Batch_size, num_patches, embed_dim)
        return x

class ViT(nn.Module):
    """
    Vision Transformer 核心模型
    """
    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__()
        
        # 1. 图像分块与嵌入层
        self.patch_embed = PatchEmbedding(in_channels, patch_size, embed_dim, img_size)
        
        # 2. CLS Token (用于最终分类的特殊 Token)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        
        # 3. 位置编码 (Positional Embedding)，长度为 num_patches + 1 (加了CLS token)
        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=dropout)
        
        # 4. Transformer 编码器
        # ViT 使用的是 Pre-Norm（即 norm_first=True），并在 MLP 层使用 GELU 激活函数
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim, 
            nhead=num_heads, 
            dim_feedforward=int(embed_dim * mlp_ratio),
            dropout=dropout, 
            activation="gelu",
            batch_first=True, 
            norm_first=True
        )
        self.blocks = nn.TransformerEncoder(encoder_layer, num_layers=depth)
        self.norm = nn.LayerNorm(embed_dim)
        
        # 5. 分类头 (Classification Head)
        self.head = nn.Linear(embed_dim, num_classes)

        # 初始化权重 (简单的正态分布初始化)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, x):
        B = x.shape[0]
        
        # 获取 Patch 嵌入: (B, num_patches, embed_dim)
        x = self.patch_embed(x)
        
        # 拼接 CLS Token: (B, 1+num_patches, embed_dim)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        
        # 加入位置编码
        x = x + self.pos_embed
        x = self.pos_drop(x)
        
        # 通过 Transformer 编码器层
        x = self.blocks(x)
        
        # 提取 CLS Token 作为整张图像的全局特征
        cls_out = x[:, 0]
        cls_out = self.norm(cls_out)
        
        # 分类预测
        out = self.head(cls_out)
        
        return out

# 测试代码
if __name__ == '__main__':
    # 模拟输入一个 Batch 的图像 (Batch_size=2, Channels=3, H=224, W=224)
    dummy_input = torch.randn(2, 3, 224, 224)
    
    # 实例化一个标准 ViT-Base 模型
    model = ViT(
        img_size=224, 
        patch_size=16,      # 常用的 16x16 
        in_channels=3, 
        num_classes=10,     # 假设是个 10 分类任务
        embed_dim=768,      # 向量维度
        depth=12,           # Transformer block 层数
        num_heads=12        # 多头注意力的头数
    )
    
    # 前向传播
    output = model(dummy_input)
    
    print(f"输入尺寸: {dummy_input.shape}")
    print(f"输出尺寸: {output.shape}") 
    # 预期输出尺寸应该是 [2, 10]，对应 2个样本，10个类别的得分