Vision Transformer (ViT) 实战教程：手把手带你训练自己的数据集（附完整源码+数据集）

前言

2020 年，谷歌发布了 Vision Transformer (ViT)，彻底打破了卷积神经网络（CNN）在计算机视觉领域的统治。它的核心思想极其暴力美学： Transformer 在 NLP 里处理文本那么强，那我直接把一张图片切成一个个小方块（Patch），当成“句子”里的“单词”喂给 Transformer 不就行了？

事实证明，大力出奇掌。Transformer 缺乏⼀些卷积神经网络固有的归纳偏置，例如平移等变性和局部性，因此在训练数据不⾜时泛化能⼒较差。然⽽，如果模型在更⼤的数据集（1400万到3亿图⽚）上进⾏训练，情况就会改变。作者发现⼤规模训练优于归纳偏置。视觉Transformer（ViT）在充分预训练后，迁移到数据较少的任务时表现出色。Transformer 模型可以说是目前最火的模型之一，目前 ChatGPT、各种开源大模型低层就是使用 Transformer 构建，那么 Vision Transformer (ViT) 可以说是当前计算机视觉模型的第一梯队。Vision Transformer (ViT) 模型结构如下图所示（摘自论文）。

从上图来看：

1. 图像 -> 块切片与展平 (Patch + Flattening)：直接把一张完整的图像像切网格一样，切成固定大小的图像块（通常是 16×16 像素）。然后，把这些二维的小块直接“拉平”成一维的向量。

2. 线性投影 (Linear Projection of Flattened Patches)：把拉平后的向量通过一个简单的线性变换（相当于全连接层），映射到一个统一的特征维度。意思是把图像块变成了 Transformer 能够处理的单词（Token）。

3. 引入类别令牌与位置编码 (Class Token & Position Embedding)：

   • 类别令牌：图中的0*块（Extra learnable [class] embedding），是一个额外加入的、可学习的特殊向量。它的唯一任务就是跟着所有图像块一起在网络里流转，最终负责汇总全局信息来输出分类结果。
   • 位置编码：为了让模型知道每个小块原本在图片上的物理位置，网络给每个块都加上了带序号的位置编码向量（图中的 0 到 9 编号）。

4. Transformer 编码器 (Transformer Encoder)：图右侧的核心组件。与 Swin 分阶段、下采样的设计不同，ViT 所有的向量打包送入编码器，编码器内部由多个相同的层堆叠而成（图中标记为 L×）。它主要包含：

   • 多头注意力机制 (Multi-Head Attention)：所有的 Patch 之间都要互相计算关联度，没有任何局部窗口的限制。
   • 多层感知机 (MLP)与层归一化 (Norm)：结合残差连接（+），进行非线性特征变换和稳定训练。

5. 分类输出 (MLP Head)：经过 L 层 Transformer 编码器的深度计算后，丢弃其他的图像块特征，只提取那个特殊的[class]向量（也就是对应0*最左边那条向上的线）的输出结果，送入最后的 MLP 模块，得出最终的分类预测（如鸟、球、车等）。

论文链接: https://arxiv.org/pdf/2010.11929
代码地址： https://github.com/google-research/vision_transformer

本篇文章将带你学习模型构建、训练自己数据集、数据集处理、推理和验证等。本篇教程适配各种分类识别任务，可以说是一个通用的教程，只需修改自己的数据集训练自己的模型即可完成自己领域的任务和系统实现。

⚡⚡常见图像分类的应用场景

可以更换自己的数据集进行训练，用 AI 列出分类模型中常见的任务如下：

1. 眼科疾病：基于深度学习的OCT图像视网膜黄斑病变多细粒度分类识别
2. 病理切片：复杂背景下的乳腺癌组织病理学显微图像良恶性分类识别
3. 血液筛查：基于深度学习的血涂片白细胞亚型高精度图像分类识别
4. 口腔健康：基于内窥镜图像的牙周炎与龋齿病变等级细粒度分类识别
5. 中医舌诊：基于深度学习的自然光照下舌象特征（齿痕、苔色）多标签分类识别
6. 皮肤病变：多中心临床数据下的罕见皮肤病种图像跨域分类识别
7. 心血管筛查：基于超声图像的心脏瓣膜结构异常评估与分类识别
8. 脑部诊断：基于深度学习的脑电图时频转换图像癫痫状态分类识别
9. 骨科辅助：X光影像下的骨折类型与骨龄阶段自动化分类识别
10. 消化道筛查：胶囊内窥镜图像下的肠道溃疡与息肉病变类型分类识别
11. 细胞分析：基于深度学习的宫颈液基细胞学涂片异常细胞分类识别
12. 遗传病筛查：基于小样本学习的面部特征遗传综合征图像分类识别
13. 寄生虫检测：基于显微图像的疟原虫感染阶段细粒度分类识别
14. 烧伤评估：基于深度学习的皮肤烧伤深度与愈合阶段图像分类识别
15. 甲状腺筛查：基于深度学习的甲状腺结节TI-RADS等级超声图像分类识别
16. 林业害虫：诱虫灯场景下的小样本林业害虫种类细粒度分类识别
17. 水产养殖：基于深度学习的密集鱼群健康状态与种类声呐图像分类识别
18. 中药鉴定：基于宏观纹理与显微特征的中药材真伪与切片分类识别
19. 杂草识别：复杂农田背景下的伴生杂草与农作物幼苗图像分类识别
20. 水果分级：基于深度学习的传送带苹果/柑橘外观瑕疵与糖度等级分类识别
21. 茶叶品控：基于微距图像的茶叶发酵程度与品种自动化分类识别
22. 土壤分析：基于深度学习的土壤剖面图像肥力与质地等级分类识别
23. 濒危动物：红外触发相机下的野生猫科动物（雪豹、华南虎）个体分类识别
24. 鸟类生态：基于深度学习的无人机视角下湿地水鸟群落种类分类识别
25. 珊瑚礁监测：水下浑浊环境下的珊瑚白化状态与种类图像分类识别
26. 种子繁育：基于高光谱图像的玉米/水稻种子活力与品种无损分类识别
27. 蘑菇采摘：自然环境下的野生可食用与有毒菌类图像高精度分类识别
28. 畜牧疫病：基于深度学习的生猪体表红斑与病变类型分类识别
29. 木材鉴定：基于显微构造图像的珍贵红木树种细粒度分类识别
30. 花粉过敏：基于气传花粉显微图像的致敏花粉种类实时分类识别
31. 新能源电池：基于深度学习的锂电池表面缺陷与极片瑕疵分类识别
32. 半导体制造：晶圆表面微米级划痕与颗粒缺陷高精度图像分类识别
33. 钢铁冶金：基于深度学习的热轧钢板表面缺陷（划痕、氧化皮）多标签分类识别
34. 纺织印染：复杂提花面料下的细微断经、断纬瑕疵类型分类识别
35. 汽车制造：基于深度学习的汽车漆面喷涂缺陷（流挂、橘皮）分类识别
36. 航空航天：飞机发动机叶片表面裂纹与损伤程度X光图像分类识别
37. 光伏产业：基于电致发光(EL)图像的太阳能电池板隐裂与黑片分类识别
38. 机械加工：基于显微视觉的数控刀具磨损状态等级分类识别
39. 焊接质检：基于深度学习的焊缝X射线探伤气孔与夹渣缺陷分类识别
40. 包装印刷：高速流水线上的药品包装字符缺失与图案错位分类识别
41. 线缆制造：基于深度学习的高压电缆横截面绝缘层气泡与杂质分类识别
42. 精密轴承：基于机器视觉的轴承滚子表面压痕与锈蚀缺陷分类识别
43. 注塑工艺：基于深度学习的塑料件飞边、缺胶等成型不良分类识别
44. 3C电子：智能手机屏幕玻璃盖板划伤与白点缺陷分类识别
45. 五金铸造：基于深度学习的压铸件内部缩孔与疏松CT图像分类识别
46. 垃圾回收：基于深度学习的工业流水线建筑垃圾（砖块、木材）分类识别
47. 气象云图：基于卫星遥感图像的台风等级与对流云团类型分类识别
48. 水质监测：基于显微暗视野图像的有害赤潮藻类细胞分类识别
49. 地质灾害：无人机视角下的滑坡、泥石流灾害类型与破坏等级分类识别
50. 极地科考：基于SAR图像的海冰类型（一年冰、多年冰）与融池分类识别
51. 城市治理：基于街景图像的城市违章建筑与市容违规现象分类识别
52. 海洋油污：基于深度学习的机载合成孔径雷达海面溢油类型分类识别
53. 岩心分析：基于高清扫描图像的深地钻探岩心岩性与沉积构造分类识别
54. 植被覆盖：基于多光谱遥感图像的荒漠化演变与植被类型分类识别
55. 光污染监测：基于夜光遥感图像的城市功能区灯光类型分类识别
56. 塑料微粒：基于拉曼光谱成像的海洋微塑料成分与形态分类识别
57. 古生物学：基于深度学习的微体古生物化石（如虫筳）种属分类识别
58. 雪崩预警：基于高山固定摄像头图像的积雪状态与雪崩风险等级分类识别
59. 大气污染：基于天空全景图像的雾霾等级与能见度分类识别
60. 湿地演化：基于多源时序遥感影像的滩涂覆被物分类识别
61. 违禁品安检：基于深度学习的复杂重叠X光图像违禁刀具与液体分类识别
62. 交通标志：雨雪天气与光照突变环境下的多国交通标志跨域分类识别
63. 车辆重识别：基于细粒度特征的套牌车与相似车型细微差异分类识别
64. 安全帽穿戴：基于深度学习的建筑工地多颜色安全帽规范穿戴分类识别
65. 特种作业：高处作业人员安全带悬挂状态与规范佩戴类型分类识别
66. 消防救援：基于深度学习的火灾初期烟雾颜色与火焰形态预警分类识别
67. 行为合规：加油站监控视角下的人员抽烟与打电话违规截帧分类识别
68. 轨道交通：基于深度学习的高铁受电弓磨损状态与异物附着分类识别
69. 无人机反制：基于雷达微多普勒频谱图的“低慢小”无人机型号分类识别
70. 伪装识别：基于深度学习的复杂丛林环境迷彩伪装人员目标分类识别
71. 路面巡检：车载视角下的大面积路面裂缝、坑槽与沉陷病害分类识别
72. 港口物流：基于深度学习的集装箱箱体残损类型与危险品标识分类识别
73. 证件防伪：基于多光谱图像的护照/身份证伪造痕迹与防伪特征分类识别
74. 煤矿安全：基于深度学习的井下防爆门状态与传送带跑偏故障分类识别
75. 电网巡检：无人机巡线视角下的绝缘子破裂与鸟巢违规搭设分类识别
76. 自助结算：基于深度学习的密集遮挡托盘菜品与烘焙糕点细粒度分类识别
77. 盲盒分拣：基于X光透视图像的盲盒内部潮玩款式分类识别
78. 服饰电商：基于弱监督学习的快时尚服装款式、领型与袖型多标签分类识别
79. 智能冰箱：基于深度学习的冰箱内部食材类别与新鲜度衰减等级分类识别
80. 文物数字化：基于深度学习的历代陶瓷器型与青花瓷纹饰断代分类识别
81. 古画鉴定：基于笔墨微观特征的中国历代水墨画流派分类识别
82. 美妆推荐：基于深度学习的素颜面部肤质（油性、干性、敏感肌）分类识别
83. 垃圾短信：基于图像文字截屏的黑灰产广告图片类型分类识别
84. 盲人辅助：基于穿戴式相机的日常面额纸币与常见障碍物细粒度分类识别
85. 宠物社区：基于深度学习的猫狗品种细粒度特征与面部状态分类识别
86. 球鞋鉴定：基于细微图像特征的限量版球鞋真伪（走线、鞋标）分类识别
87. 珠宝估值：基于显微图像的钻石净度等级与翡翠种水高精度分类识别
88. 智能书柜：基于书脊图像的图书馆错架图书与索书号匹配分类识别
89. 字体设计：基于深度学习的残损古代碑刻书法字体风格分类识别
90. 动漫创作：基于手绘草图的二次元角色发型与服装部件分类识别
91. 无人售货：基于深度学习的开放式货架商品乱放与错位状态分类识别
92. 酒类品鉴：基于包装细节的进口红酒年份、产区与真伪标签分类识别
93. 中式家具：基于木纹理与榫卯结构的明清古典家具制式分类识别
94. 票据自动化：基于深度学习的复杂扭曲医疗发票与增值税专票板式分类识别
95. 外卖审核：基于深度学习的骑手送达照片规范性（门牌、餐品）分类识别
96. 乐器维修：基于声学频谱图转视觉图像的提琴面板共振缺陷分类识别
97. 健身指导：基于骨骼点热图图像的用户深蹲、硬拉动作标准程度分类识别
98. 植物园艺：基于深度学习的多肉植物品种与缺水/徒长状态等级分类识别
99. 硬币收藏：基于磨损图像的古钱币/纪念币品相评级与版别分类识别
100. 星空摄影：基于深度学习的深空天体（星云、星系、星团）形态学分类识别

🐴一、数据集准备

获取数据集，永远先问一句：“有没有现成的？”，开源社区有大量数据集，使用现成的数据集可以节省 70% 时间，如果没有数据集可以自己选择爬虫。下面推荐获取深度学习数据集的网站，仅供学习。

⚡⚡社区下载

☑️1.1 国外网站

1. Kaggle

• 网址：www.kaggle.com/datasets
• 特点：全球最大的数据科学社区之一。
• 技巧：搜索关键词如 “Cat and Dog”, “Flower Classification”。
• 下载：注册账号后直接下载，或者使用 Kaggle API 命令行下载。

2. Hugging Face Datasets

• 网址：huggingface.co/datasets