1. 为什么要在 Keras 里手搓一个 Vision TransformerTransformer 架构从自然语言处理领域杀进计算机视觉最标志性的事件就是 Vision TransformerViT的提出。它干了一件很“暴力”的事把一张图切成固定大小的小块拉平之后当成词序列直接丢给标准 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×3patch 大小是 16×16那么 patch 数量就是 (224/16) × (224/16) 14 × 14 196 个。每个 patch 展平后是 16×16×3 768 维正好映射到 D768这一步在很多实现里干脆用一个步长等于 patch 大小的卷积来完成既高效又等价。这里有个细节值得说为什么用卷积实现 Patch Embedding因为卷积核大小和步长都设成 16输出就是每个 patch 的线性投影数学上和“切块展平全连接”完全一致但计算效率高得多还能顺便利用框架对卷积的优化。这是我在实际写代码时首选的方式代码量少跑得还快。2.2 类别 token 与位置编码两个必须加的“配件”光有 patch 序列还不够。Transformer 本身对序列顺序是无感知的而图像 patch 之间的空间关系很重要所以必须加位置编码。ViT 用的是可学习的位置编码也就是一个形状为 (N1, D) 的矩阵直接和输入序列相加。为什么是可学习而不是正弦固定编码原论文做过对比两者效果接近但可学习编码更简单也方便迁移到不同分辨率虽然插值会麻烦一点。另一个配件是类别 tokenclass 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 继承下来的经验值实验上效果和计算量的平衡比较好。多头注意力的头数一般设为 12D768 时每个头的维度是 64。头数太少表达能力受限头数太多单头维度太小也不好。这个比例在实际调参时可以作为起点。3. Keras 复现的核心模块拆解3.1 多头自注意力的 Keras 实现要点多头自注意力是整个 ViT 里最绕的部分但拆开看并不复杂。输入序列形状是 (B, N1, D)先经过三个全连接层分别得到 Q、K、V形状都是 (B, N1, D)。然后按头数 h 把 D 拆成 h 份每份维度 d D/h变成 (B, h, N1, d)。接着计算注意力分数Q 和 K 的转置做矩阵乘除以 sqrt(d) 缩放再经过 softmax最后和 V 相乘得到 (B, h, N1, 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 位置编码的初始化与插值位置编码矩阵的形状是 (N1, D)用Embedding层或者直接AddWeight都可以。初始化一般用截断正态分布标准差 0.02。训练时它和模型其他参数一起更新。如果预训练模型的分辨率和你的输入分辨率不一致就需要对位置编码做插值。常见做法是把 (N1, 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×32patch 大小可以设成 4×4这样 patch 数量是 64序列长度适中训练快。如果想贴近原论文用 224×224 输入、16×16 patch但那样需要更大显存和更多数据。4.2 模型定义逐层搭建 ViT下面是我实际用的模型定义骨架参数以 CIFAR-10 为例输入 32×32×3patch 大小 4D128深度 6头数 8MLP 维度 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_sizepatch_size, stridespatch_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 xPatch 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), initializertruncated_normal, trainableTrue) self.pos_embed self.add_weight(shape(1, num_patches 1, embed_dim), initializertruncated_normal, trainableTrue) 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], axis1) return x self.pos_embedTransformer Block 部分class TransformerBlock(layers.Layer): def __init__(self, embed_dim, num_heads, mlp_dim, dropout0.1): super().__init__() self.norm1 layers.LayerNormalization(epsilon1e-6) self.attn layers.MultiHeadAttention(num_headsnum_heads, key_dimembed_dim // num_heads, dropoutdropout) self.norm2 layers.LayerNormalization(epsilon1e-6) self.mlp tf.keras.Sequential([ layers.Dense(mlp_dim, activationtf.keras.activations.gelu), layers.Dropout(dropout), layers.Dense(embed_dim), layers.Dropout(dropout) ]) def call(self, x, trainingFalse): x x self.attn(self.norm1(x), self.norm1(x), trainingtraining) x x self.mlp(self.norm2(x), trainingtraining) return x最后组装模型取 class token 输出接分类头def build_vit(input_shape, patch_size, embed_dim, depth, num_heads, mlp_dim, num_classes): inputs layers.Input(shapeinput_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(epsilon1e-6)(x) cls_output x[:, 0] outputs layers.Dense(num_classes, activationsoftmax)(cls_output) return Model(inputs, outputs)4.3 训练配置与超参数选择编译时用 AdamW 优化器学习率 3e-4权重衰减 0.05batch size 64训练 100 个 epoch。损失函数用稀疏交叉熵。学习率调度用余弦退火前 5 个 epoch 做线性 warmup。为什么用 AdamW 而不是 Adam因为 AdamW 把权重衰减和梯度更新解耦正则效果更稳定Transformer 类模型基本都用它。optimizer tf.keras.optimizers.AdamW(learning_rate3e-4, weight_decay0.05) model.compile(optimizeroptimizer, losssparse_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×4num_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 时 N197注意力矩阵是 197×197还好。但如果 patch 设成 8×8N 变成 785显存直接翻好几倍。优化手段有几个一是用混合精度训练tf.keras.mixed_precision.set_global_policy(mixed_float16)显存能省三成左右二是减小 batch size三是用梯度累积模拟大 batch。推理速度方面Keras 的MultiHeadAttention在 GPU 上已经优化得不错但如果你的部署环境是 CPU建议把模型转成 TFLite 或 ONNX速度会快很多。我实测下来同样一个 6 层 ViTTFLite 在移动端 CPU 上比原生 Keras 快 2 到 3 倍。5.4 常见问题速查表问题现象可能原因排查与解决loss 变 NaN学习率过大、LayerNorm epsilon 过小降学习率到 1e-4epsilon 改 1e-6验证准确率不涨数据量不足、缺少增强加 RandAugment、MixUp或用预训练权重形状报错num_patches 算错、key_dim 不匹配打印各张量形状断言 embed_dim % num_heads 0显存溢出序列过长、batch 过大混合精度、减 batch、增大 patch训练极慢未用 GPU、注意力实现低效确认 GPU 可用用内置 MultiHeadAttention6. 迁移学习与后续扩展方向6.1 加载预训练权重的正确姿势自己从零训 ViT在小数据集上很难达到理想精度。更实际的做法是加载在大规模数据集上预训练好的权重然后微调。Keras 里没有官方 ViT 权重但你可以从其他框架转换或者用tf.keras.applications里有没有对应模型目前没有需要自己转。转换时注意位置编码的插值如果预训练分辨率是 224你的输入是 32位置编码要下采样否则序列长度对不上。微调时学习率要调小比如 1e-5并且冻结前面几层只训后面几层和分类头。这样收敛快也不容易过拟合。我试过在 CIFAR-10 上微调一个预训练 ViT10 个 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会发现它们无非是在这些基础模块上做排列组合和优化。基础打牢了后面学什么都快。