参考DINOv3源码改进手撕ViT模型的目标检测性能

在本公众号前一段时间关于ViT的连载连载6：手撕一个Vision Transformer（ViT）目标检测模型（模型推理篇）等中，我从0开始编写了一个ViT架构的模型。

从应用效果上看，在图像分类方面，模型表现得相当不错，而在做目标检测任务时却表现不佳，这一方面是因为YOLO128数据集图像太少，另一方面模型也有一些改进空间。

本篇根据DINOv3模型源码里的一些思路，对手撕的ViT稍做调整，以改善其在目标检测方面的性能。具体来说，一是找更大的数据集进行训练测试；二是改进嵌入层；三是调整模型的自注意力头个数等参数。

01 数据集

我们找一个稍大的“非洲野生动物数据集”，共有1504张图像，4个分类（水牛、大象、犀牛和斑马）。数据集的详细说明在：

https://docs.ultralytics.com/zh/datasets/detect/african-wildlife/#applications

下载地址为：

https://github.com/ultralytics/assets/releases/download/v0.0.0/african-wildlife.zip

将下载的压缩包解压替换上一个连载中的coco128_yolo目录下的文件，为了减少代码的修改，可将train、val目录改为train2017、val2017，这样原来装载YOLO128的代码就不用做任何修改了。

02 Patch嵌入

遵从ViT的论文，在从RGB数据变换成嵌入向量时，我用了nn.Linear线性层展平，之前也看了一些文章，这个环节多用卷积网络处理，但我不想在Transformer中再引入卷积网络，意在做一个“纯”的ViT。

但在看了DINOv3的实现源码后，发现自己的想法草率了，DINOv3人家也用卷积。具体见：

https://github.com/facebookresearch/dinov3/blob/main/dinov3/layers/patch_embed.py

所以我也将原来的nn.Linear换成nn.Conv2d，卷积核等于Patch的大小（16）时，他的结果与Patch切分是等效的，而且更快，在后续的训练过程中也证明了这一点。

原来的模型用了8个注意力头，嵌入维度为768，研究了一下DINOv3模型定义的源码：

https://github.com/facebookresearch/dinov3/blob/main/dinov3/models/vision_transformer.py

按DINOv3 vit small的规模，把手撕的模型的自注意力头从8个减成6个，嵌入维度减半到384，这样训练更快，模型更小。

调整后的模型定义、训练和推理代码名为vit_detr_model1.py、vit_detr_train1.py和vit_detr_pred1.py，也在原仓库下：

https://github.com/little51/pure-vit-examples

经过100轮训练后，从训练集、验证集和测试集各找一张图片进测试，结果如下：

从结果上看，类别判断基本正确、边框范围精度有待提高和置信度较低。

但要注意的是：这可是从0开始训练的ViT模型，并未使用如DINOv3、YOLO那些预训练模型，关于图像的特征、分类、目标检测都是从0开始学习的，达到这个情况也还算不错。

如果为了深入了解ViT的细节，可以再从位置编码、Transformer结构等方面继续改进。如果是用于生产，则没有改进的必要，直接用DINOv3预训练模型+训练任务头或微调YOLO，都会得到较好的结果，因为这些预训练模型已学习到了大量的图像特征。手撕模型再怎么改，毕竟数据集太少，效果也会一般。