最近在学习 Vision Transformer(ViT),跟着教程把模型从 patch embedding 到图像分类的流程实现了一遍,并整理成了一个项目,分享给同样在学习 ViT 的朋友。
先说明:这是个人学习笔记,并非原创教程或原创模型。 主要参考 Daniel Bourke 的 Learn PyTorch 第 08 章,在此感谢原作者。
项目包含什么?
-
使用 PyTorch 搭建 patch embedding、多头自注意力、MLP 和 Transformer Encoder。
-
为手写 ViT 加载 torchvision 的预训练主干权重,完成图片三分类。
-
整理数据加载、训练、权重保存、单张图片预测等完整流程。
为了便于理解,我把分步骤的讲解与实验放在了 main.ipynb 中,同时提供模块化脚本,方便对照阅读和单独运行。
项目地址
环境配置、文件说明、运行方式和注意事项已经写在仓库的 README 中,这里就不重复展开了。建议先阅读 README,再按顺序查看 main.ipynb,结合每一层的输入输出形状理解代码。