2. ViT核心原理

好，咱们今天来啃ViT最核心的几个概念。说实话，我第一次看ViT的时候，也被这几个名词搞得有点晕。但后来我发现，只要搞懂了这四个东西——Patch Embedding、Transformer Encoder、Class 、位置编码——ViT的骨架你就拿下了。

我个人习惯，学一个新架构，先别急着看公式，先搞清楚它要解决什么问题。ViT说白了就是：把图像切成小块，当成序列送给Transformer去处理。就这么简单。

1. Patch Embedding：把图像变成序列

CNN处理图像，天然就是二维卷积核滑来滑去。但Transformer呢？它只认识一维的序列。所以第一步，得把二维图像转成一维。

怎么转？ViT的做法很直接——把图像切成固定大小的patch。

假设输入是224x224的彩色图，patch大小设为16x16。那么一共会切出 (224/16)² = 196 个patch。每个patch有 16x16x3 = 768 个像素值。

然后，用一个线性层（全连接层）把这768个像素映射成一个向量。这个向量就是patch embedding。嗯，这里要注意：每个patch共享同一个线性层，相当于对所有patch做同样的变换。

我在项目中遇到过一个问题：patch size怎么选？太小了，序列太长，计算量爆炸；太大了，空间细节丢失。16x16是个比较折中的选择，MobileViT里用了更小的patch，因为移动端算力有限。

# PyTorch实现Patch Embedding（简化版）
import torch
import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        # 直接用卷积实现：同时完成切patch和线性映射
        self.proj = nn.Conv2d(in_chans, embed_dim, 
                              kernel_size=patch_size, stride=patch_size)
    
    def forward(self, x):
        # x: [B, 3, 224, 224]
        x = self.proj(x)  # [B, 768, 14, 14]
        x = x.flatten(2)  # [B, 768, 196]
        x = x.transpose(1, 2)  # [B, 196, 768]
        return x

2. Class ：分类的“锚点”

好，现在我们有196个patch 了。但问题来了：Transformer输出的是196个向量，到底用哪个做分类？

ViT的做法很巧妙——额外加一个可学习的class 。这个跟patch 一起送入Transformer，最后只用它对应的输出做分类。

为什么这样设计？我个人理解是：class 相当于一个“问号”，它在训练过程中学会去“询问”其他patch的信息，最终自己汇总出全局特征。你想想看，如果没有这个，就得对196个输出做平均池化，那会丢失很多信息。

# Class 的初始化
self.cls_ = nn.Parameter(torch.randn(1, 1, embed_dim))

# 前向传播时拼接到patch 前面
cls_tokens = self.cls_.expand(B, -1, -1)  # [B, 1, 768]
x = torch.cat((cls_tokens, x), dim=1)  # [B, 197, 768]

3. 位置编码：告诉模型“谁在哪儿”

Transformer本身是置换不变的——你把patch的顺序打乱，它输出的结果一样。这对序列建模来说是个大问题。图像里左上角的patch和右下角的patch，语义完全不同。

所以，我们需要给每个加上位置信息。ViT用的是可学习的位置编码，说白了就是一个跟patch数量一样大的参数矩阵，训练时一起更新。

为什么不用正弦编码？ViT里做了实验，发现可学习的位置编码效果差不多，但实现更简单。我个人在移动端部署时，更喜欢可学习的，因为推理时少算一堆三角函数。

# 位置编码初始化
self.pos_embed = nn.Parameter(torch.randn(1, num_patches + 1, embed_dim))

# 加到patch embedding上
x = x + self.pos_embed

4. Transformer Encoder：核心处理单元

有了序列，接下来就是Transformer Encoder上场了。ViT用的就是标准Transformer的Encoder部分，没有Decoder。

一个Encoder Block包含两个核心子层：

• Multi-Head Self-Attention (MSA)：让每个看到所有其他
• MLP (Feed Forward)：对每个独立做非线性变换

每个子层后面都跟着Layer Normalization和残差连接。结构很简单，但效果惊人。

我记得第一次在手机上跑ViT时，最头疼的就是Self-Attention的计算量。序列长度197（196个patch + 1个class ），每个都要跟其他所有做点积，复杂度是O(n²)。

后来MobileViT用了更轻量的注意力机制，才把模型塞进手机里。这部分咱们后面章节会详细讲。

# Transformer Encoder Block（简化版）
class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim)
        )
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        x = x + self.mlp(self.norm2(x))
        return x

5. 整体流程：一张图看懂

咱们把上面四个部分串起来，看看ViT的完整前向流程：

1. 输入图像 → 切成patch → 线性映射得到patch embedding
2. 拼接class 到序列最前面
3. 加上位置编码
4. 送入L层Transformer Encoder
5. 取出class 对应的输出 → 过分类头 → 得到预测结果

6. 移动端部署的视角

聊完原理，咱们从移动端部署的角度看看这几个组件：

组件	计算量占比	移动端优化点
Patch Embedding	低（~1%）	用Conv2d实现，NCNN/TNN直接支持
位置编码	极低（<0.1%）	可提前融合到权重中，推理时无额外开销
Self-Attention	高（~60%）	需要优化：稀疏注意力、窗口注意力等
MLP	中（~30%）	可以用卷积替代全连接，利用NCNN的卷积优化

好了，ViT的核心原理就这些。说白了就是：切patch、加class 、加位置编码、过Transformer。下一章咱们会深入MobileViT，看看它怎么在保持ViT优势的同时，把模型做得更轻量。