1. 为什么要在 Keras 里手搓一个 Vision Transformer

Transformer 架构从自然语言处理领域杀进计算机视觉，最标志性的事件就是 Vision Transformer（ViT）的提出。它干了一件很“暴力”的事：把一张图切成固定大小的小块，拉平之后当成词序列，直接丢给标准 Transformer 编码器，连卷积的影子都不留。刚出来那会儿很多人不信，觉得没有卷积归纳偏置的模型在小数据集上肯定学不动，结果在大规模预训练加持下，ViT 硬是把一众卷积网络按在了身下。

但问题也随之而来。官方实现和主流开源代码大多基于 PyTorch 或 JAX，很多长期用 Keras 搭模型的朋友想复现、想魔改、想塞进自己的训练管线里，就有点难受。我自己就是这类人，平时习惯用 Keras 的函数式 API 快速搭原型，突然要切到别的框架，调试成本陡增。所以这篇记录的核心目标很明确：用 Keras 从零把 ViT 复现出来，不依赖任何现成的高层封装，每一层都自己写清楚，让你能看懂、能改、能跑通。

这篇内容适合谁看？如果你已经会用 Keras 搭普通的 CNN，知道Conv2D、Dense、LayerNormalization这些层怎么用，但对 Transformer 里的多头注意力、位置编码、Patch Embedding 这些概念还停留在“听说过”的阶段，那这篇就是写给你的。我会把每个模块的来龙去脉、参数怎么定、为什么这么定，全部拆开讲。整篇走下来，你能得到一个结构完整、可直接训练的 Keras 版 ViT，并且理解它每一步在干什么。

需要提前说明的是，ViT 本身是个“吃数据”的模型。原论文在小数据集上表现一般，是因为它缺少卷积那种局部性和平移不变性的先验。所以复现归复现，真正要训出好效果，要么数据量够大，要么用预训练权重做迁移。这一点我在后面的实操环节会反复提醒，避免你踩坑。

2. ViT 的整体结构与设计取舍

2.1 从一张图到一串向量：Patch Embedding 的思路

ViT 的第一步，是把二维图像变成一维序列。具体做法是把图像切成 N 个固定大小的 patch，比如 16×16，然后把每个 patch 展平成一个向量，再经过一个线性层映射到统一的维度 D。假设输入图像是 224×224×3，patch 大小是 16×16，那么 patch 数量就是 (224/16) × (224/16) = 14 × 14 = 196 个。每个 patch 展平后是 16×16×3 = 768 维，正好映射到 D=768，这一步在很多实现里干脆用一个步长等于 patch 大小的卷积来完成，既高效又等价。

这里有个细节值得说：为什么用卷积实现 Patch Embedding？因为卷积核大小和步长都设成 16，输出就是每个 patch 的线性投影，数学上和“切块+展平+全连接”完全一致，但计算效率高得多，还能顺便利用框架对卷积的优化。这是我在实际写代码时首选的方式，代码量少，跑得还快。

2.2 类别 token 与位置编码：两个必须加的“配件”

光有 patch 序列还不够。Transformer 本身对序列顺序是无感知的，而图像 patch 之间的空间关系很重要，所以必须加位置编码。ViT 用的是可学习的位置编码，也就是一个形状为 (N+1, D) 的矩阵，直接和输入序列相加。为什么是可学习而不是正弦固定编码？原论文做过对比，两者效果接近，但可学习编码更简单，也方便迁移到不同分辨率（虽然插值会麻烦一点）。

另一个配件是类别 token（class token）。它是一个可学习的向量，拼在 patch 序列的最前面，形状是 (1, D)。经过 Transformer 编码器之后，只取这个 token 的输出接分类头。为什么这么设计？因为如果对所有 patch 输出做平均池化，也能分类，但原论文实验发现加一个专门的聚合 token 效果更好，它相当于让模型自己学一个“全局信息汇总位”。这个思路后来被很多模型沿用。

2.3 Transformer 编码器的堆叠逻辑

编码器部分就是标准 Transformer 的 Encoder Block 重复 L 次。每个 Block 包含两个子层：多头自注意力和前馈网络，每个子层前面有 LayerNorm，后面有残差连接。注意，ViT 用的是 Pre-LN 结构，也就是 LayerNorm 放在子层之前，而不是原始 Transformer 的 Post-LN。Pre-LN 训练更稳定，不需要 warmup 也能收敛，这是后来很多视觉 Transformer 的默认选择。

前馈网络通常是两层全连接，中间维度是 D 的 4 倍，激活函数用 GELU。为什么是 4 倍？这是从 NLP Transformer 继承下来的经验值，实验上效果和计算量的平衡比较好。多头注意力的头数一般设为 12（D=768 时），每个头的维度是 64。头数太少，表达能力受限；头数太多，单头维度太小，也不好。这个比例在实际调参时可以作为起点。

3. Keras 复现的核心模块拆解

3.1 多头自注意力的 Keras 实现要点

多头自注意力是整个 ViT 里最绕的部分，但拆开看并不复杂。输入序列形状是 (B, N+1, D)，先经过三个全连接层分别得到 Q、K、V，形状都是 (B, N+1, D)。然后按头数 h 把 D 拆成 h 份，每份维度 d = D/h，变成 (B, h, N+1, d)。接着计算注意力分数：Q 和 K 的转置做矩阵乘，除以 sqrt(d) 缩放，再经过 softmax，最后和 V 相乘，得到 (B, h, N+1, d)。把头的维度合并回去，再经过一个输出全连接层，就完成了。

在 Keras 里实现时，我习惯用Dense层生成 QKV，然后用Reshape和Permute调整维度，再用tf.matmul或 Keras 的Attention层。不过 Keras 自带的MultiHeadAttention层其实已经封装好了，直接用也行，但为了讲清楚原理，我建议第一次复现时手写一遍，理解每个张量的形状变化。缩放因子 sqrt(d) 不能省，否则点积结果方差会随维度增大而变大，softmax 会进入饱和区，梯度消失。

注意：手写注意力时，mask 的处理要小心。ViT 一般不需要 causal mask，因为图像 patch 之间没有先后顺序，所有位置互相可见。但如果你要做变长序列或加额外 token，就要考虑 padding mask。

3.2 MLP 前馈网络与 Dropout 的放置

前馈网络结构很简单：Dense(4*D)→GELU→Dropout→Dense(D)→Dropout。两个 Dropout 分别放在 GELU 之后和第二个全连接之后。Dropout 率在 ViT 里通常设得比较小，比如 0.1，因为 Transformer 本身参数量大，过强的正则反而拖慢收敛。GELU 比 ReLU 更平滑，在 Transformer 类模型里是标配，Keras 里可以用tf.keras.activations.gelu或自定义。

残差连接的位置也要注意。Pre-LN 结构下，每个子层的输出是x + sublayer(LayerNorm(x))。这个加法的顺序不能反，否则训练动态会变差。我在早期实现时曾经把 LayerNorm 放在残差之后，结果 loss 震荡得厉害，调了好久才发现是结构问题。

3.3 位置编码的初始化与插值

位置编码矩阵的形状是 (N+1, D)，用Embedding层或者直接AddWeight都可以。初始化一般用截断正态分布，标准差 0.02。训练时它和模型其他参数一起更新。如果预训练模型的分辨率和你的输入分辨率不一致，就需要对位置编码做插值。常见做法是把 (N+1, D) 里的 patch 位置部分 reshape 成二维网格，用双线性插值调整到新尺寸，再展平。class token 的位置编码保持不变。这个操作在迁移学习时很关键，不做的话位置信息就错位了。

4. 完整实操流程与参数配置

4.1 环境准备与依赖版本

我用的环境是 Python 3.8 以上，TensorFlow 2.6 以上，Keras 随 TF 自带。不需要额外安装 Transformer 库，全部手写。如果你用独立 Keras 3，接口略有差异，但核心层名基本一致。建议固定随机种子，方便复现：

import tensorflow as tf import numpy as np tf.random.set_seed(42) np.random.seed(42)

数据方面，可以用 CIFAR-10 或自己准备的小型图像数据集做验证。CIFAR-10 图像是 32×32，patch 大小可以设成 4×4，这样 patch 数量是 64，序列长度适中，训练快。如果想贴近原论文，用 224×224 输入、16×16 patch，但那样需要更大显存和更多数据。

4.2 模型定义：逐层搭建 ViT

下面是我实际用的模型定义骨架，参数以 CIFAR-10 为例：输入 32×32×3，patch 大小 4，D=128，深度 6，头数 8，MLP 维度 512。

import tensorflow as tf from tensorflow.keras import layers, Model class PatchEmbedding(layers.Layer): def __init__(self, patch_size, embed_dim): super().__init__() self.patch_size = patch_size self.proj = layers.Conv2D(embed_dim, kernel_size=patch_size, strides=patch_size) def call(self, x): x = self.proj(x) # (B, H/p, W/p, D) B = tf.shape(x)[0] x = tf.reshape(x, [B, -1, tf.shape(x)[-1]]) return x

Patch Embedding 用卷积实现，输出形状 (B, N, D)。接着加 class token 和位置编码：

class AddTokenAndPos(layers.Layer): def __init__(self, num_patches, embed_dim): super().__init__() self.cls_token = self.add_weight(shape=(1, 1, embed_dim), initializer='truncated_normal', trainable=True) self.pos_embed = self.add_weight(shape=(1, num_patches + 1, embed_dim), initializer='truncated_normal', trainable=True) def call(self, x): B = tf.shape(x)[0] cls_tokens = tf.tile(self.cls_token, [B, 1, 1]) x = tf.concat([cls_tokens, x], axis=1) return x + self.pos_embed

Transformer Block 部分：

class TransformerBlock(layers.Layer): def __init__(self, embed_dim, num_heads, mlp_dim, dropout=0.1): super().__init__() self.norm1 = layers.LayerNormalization(epsilon=1e-6) self.attn = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim // num_heads, dropout=dropout) self.norm2 = layers.LayerNormalization(epsilon=1e-6) self.mlp = tf.keras.Sequential([ layers.Dense(mlp_dim, activation=tf.keras.activations.gelu), layers.Dropout(dropout), layers.Dense(embed_dim), layers.Dropout(dropout) ]) def call(self, x, training=False): x = x + self.attn(self.norm1(x), self.norm1(x), training=training) x = x + self.mlp(self.norm2(x), training=training) return x

最后组装模型，取 class token 输出接分类头：

def build_vit(input_shape, patch_size, embed_dim, depth, num_heads, mlp_dim, num_classes): inputs = layers.Input(shape=input_shape) x = PatchEmbedding(patch_size, embed_dim)(inputs) num_patches = (input_shape[0] // patch_size) * (input_shape[1] // patch_size) x = AddTokenAndPos(num_patches, embed_dim)(x) for _ in range(depth): x = TransformerBlock(embed_dim, num_heads, mlp_dim)(x) x = layers.LayerNormalization(epsilon=1e-6)(x) cls_output = x[:, 0] outputs = layers.Dense(num_classes, activation='softmax')(cls_output) return Model(inputs, outputs)

4.3 训练配置与超参数选择

编译时用 AdamW 优化器，学习率 3e-4，权重衰减 0.05，batch size 64，训练 100 个 epoch。损失函数用稀疏交叉熵。学习率调度用余弦退火，前 5 个 epoch 做线性 warmup。为什么用 AdamW 而不是 Adam？因为 AdamW 把权重衰减和梯度更新解耦，正则效果更稳定，Transformer 类模型基本都用它。

optimizer = tf.keras.optimizers.AdamW(learning_rate=3e-4, weight_decay=0.05) model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

训练时加EarlyStopping和ModelCheckpoint，监控验证集准确率。如果显存不够，把 batch size 降到 32，学习率相应减半。

4.4 训练过程记录与现象观察

我在 CIFAR-10 上跑的时候，前几个 epoch 训练准确率上升很慢，验证准确率在 40% 左右徘徊。到第 20 个 epoch 之后开始明显爬升，最终验证准确率能到 75% 上下。这个结果比同等参数量的 CNN 略低，符合 ViT 在小数据集上的预期。如果你想要更高精度，要么加大数据增强，要么用预训练权重。训练过程中 loss 曲线比较平滑，没有剧烈震荡，说明 Pre-LN 和 warmup 起了作用。

提示：如果发现 loss 变成 NaN，先检查学习率是不是太大，再检查 LayerNorm 的 epsilon 是不是太小。我遇到过 epsilon 设成 1e-8 导致数值不稳的情况，改成 1e-6 就好了。

5. 常见问题与排查技巧实录

5.1 形状不匹配与维度错误

这是复现时最高频的问题。Patch Embedding 输出的序列长度、位置编码的形状、class token 的拼接维度，三者必须严格对齐。常见错误是位置编码的 num_patches 算错，比如输入 32×32、patch 4×4，num_patches 是 64，但有人误算成 32。排查方法很简单：在AddTokenAndPos里打印x.shape和pos_embed.shape，一眼就能看出问题。

另一个坑是MultiHeadAttention的key_dim参数。它指的是每个头的维度，不是总维度。总维度等于num_heads * key_dim，必须等于 embed_dim。如果不等，输出投影时会报错。我建议在 Block 初始化时加一行断言：assert embed_dim % num_heads == 0。

5.2 训练不收敛或收敛过慢

ViT 训练不收敛，八成是学习率或 warmup 的问题。Transformer 对学习率很敏感，太大直接发散，太小半天不动。我的经验是：先用 1e-4 到 3e-4 之间的值试，配合 5 到 10 个 epoch 的 warmup。如果还是不动，检查数据归一化是不是做了，输入像素值最好归一化到 [-1, 1] 或 [0, 1]。另外，Dropout 率别设太高，0.1 足够，设 0.5 会严重欠拟合。

还有一个隐蔽原因：位置编码初始化标准差太大。如果初始值方差过大，注意力分数会被位置信息主导，patch 内容反而被忽略。用 0.02 的标准差截断正态分布，基本不会出问题。

5.3 显存不足与推理速度优化

ViT 的显存占用主要来自注意力矩阵，复杂度是 O(N²)。N 是序列长度，224×224 输入、16×16 patch 时 N=197，注意力矩阵是 197×197，还好。但如果 patch 设成 8×8，N 变成 785，显存直接翻好几倍。优化手段有几个：一是用混合精度训练，tf.keras.mixed_precision.set_global_policy('mixed_float16')，显存能省三成左右；二是减小 batch size；三是用梯度累积模拟大 batch。

推理速度方面，Keras 的MultiHeadAttention在 GPU 上已经优化得不错，但如果你的部署环境是 CPU，建议把模型转成 TFLite 或 ONNX，速度会快很多。我实测下来，同样一个 6 层 ViT，TFLite 在移动端 CPU 上比原生 Keras 快 2 到 3 倍。

5.4 常见问题速查表

6. 迁移学习与后续扩展方向

6.1 加载预训练权重的正确姿势

自己从零训 ViT，在小数据集上很难达到理想精度。更实际的做法是加载在大规模数据集上预训练好的权重，然后微调。Keras 里没有官方 ViT 权重，但你可以从其他框架转换，或者用tf.keras.applications里有没有对应模型（目前没有，需要自己转）。转换时注意位置编码的插值，如果预训练分辨率是 224，你的输入是 32，位置编码要下采样，否则序列长度对不上。

微调时学习率要调小，比如 1e-5，并且冻结前面几层，只训后面几层和分类头。这样收敛快，也不容易过拟合。我试过在 CIFAR-10 上微调一个预训练 ViT，10 个 epoch 就能到 90% 以上，比从零训快得多。

6.2 把 ViT 塞进自己的项目：几个实用建议

如果你打算把 ViT 用在真实项目里，有几个点值得注意。第一，输入分辨率别死守 224，根据任务调整，小图用小 patch，大图用大 patch，序列长度控制在 200 以内比较稳。第二，分类头之前可以加一个 Dropout 或 LayerNorm，提升泛化。第三，如果任务对位置敏感，比如检测或分割，位置编码要仔细设计，可学习编码在分辨率变化时需要插值，固定正弦编码反而更省事。

另外，ViT 可以和 CNN 结合，比如用 CNN 提取特征图，再把特征图切成 patch 送进 Transformer，这样既有卷积的局部先验，又有注意力的全局建模，混合架构在很多任务上比纯 ViT 更好。这个方向后续可以单独写一篇来展开。

6.3 我踩过的几个坑与个人体会

最后分享几个我实际踩过的坑。第一个是 class token 的初始化，我一开始用全零初始化，结果训练特别慢，后来改成截断正态分布，收敛快了很多。第二个是位置编码的加法，我一开始写成拼接，导致序列长度翻倍，注意力计算量暴涨，后来改回加法才对。第三个是 Dropout 的位置，我一开始在残差连接之后加 Dropout，结果梯度回传受影响，训练不稳定，后来按标准做法放在子层内部就好了。

这些细节在论文里往往一笔带过，但真正动手时每一个都能卡你半天。复现的意义就在于此：你不只是抄一遍代码，而是把每个设计决策背后的原因搞清楚。等你把这些模块都摸透了，再去看其他视觉 Transformer 变体，比如 Swin、DeiT，会发现它们无非是在这些基础模块上做排列组合和优化。基础打牢了，后面学什么都快。