Vision Transformer Pytorch 实现代码学习记录

雪天鱼

已于 2024-03-01 20:27:23 修改

阅读量1.1k

点赞数 23

分类专栏：深度学习文章标签： transformer pytorch 学习

于 2024-02-13 22:10:41 首次发布

本文链接：https://blog.youkuaiyun.com/qq_44447544/article/details/136110021

版权

深度学习专栏收录该内容

28 篇文章

订阅专栏

目前运营的社交平台账号：

优快云【雪天鱼】: 雪天鱼-优快云博客
哔哩哔哩【雪天鱼】: 雪天鱼个人主页-bilibili.com

可能后续有更新，也可能没有更新，谨慎参考

V1.0 24-02-13 ViT 代码的基本训练, 预测推理脚本运行

1 学习目标

能用官方的 ViT 预训练模型在 imagenet1k 上进行预测推理完成
在 ImageNet-1K 的完整验证集上验证下载的官方 ViT 预训练模型的准确率

未处理的问题：

官方的 ViT 预处理模型训练时的图片数据预处理方法是什么？

Github pytorch实现的 ViT 代码下载：deep-learning-for-image-processing/pytorch_classification/vision_transformer at master · WZMIAOMIAO/deep-learning-for-image-processing · GitHub
Note: 非官方仓库代码，但 vit_model.py 即ViT 模型定义代码是用的被 TIMM 采用的代码。

已经处理好的 ImageNet1K数据集网盘链接：
链接：https://pan.baidu.com/s/1sYMIwqkNldmqpaJqDK8lSQ?pwd=2024
提取码：2024

2 运行 `flops.py` (不重要，可跳过)

先安装fvcore包： pip install fvcore
然后点击运行会出错，报错为:
ValueError: Invalid type <class 'numpy.int32'> for the flop count! Please use a wider type to avoid overflow.

|850

点击红框中的位置进入到 jit_handles.py 文件中，修改 14~19行代码如下：

try:  
    from math import prod  
except ImportError:  
    from numpy import prod as prodnp  
    def prod(x):   
        return int(prodnp(x))

然后再重新运行 flops.py 无报错。结果为：

Self-Attention FLOPs: 60129542144
# 中间有一些红色字体的 warnings
Multi-Head Attention FLOPs: 68719476736

3 训练—train.py

从 vit_model 中导入想要训练的 ViT版本, 把默认导入的 vit_base_patch16_224_in21k 给注释掉，确保加载的预训练权重和实例化的模型class一致。

from vit_model import vit_base_patch16_224 as create_model

运行脚本，默认训练10 epochs, 每轮都会将训练好的权重文件保存至 weights 目录下

模型有 327 MB
用tensorboard 打开 runs 目录下的训练log，如下图所示：

4 预测推理—predict.py

现在我们用训练好的模型进行预测推理，自己从数据集或者网上选择一张图作为输入，预测结果如下图所示：

5 在 ImageNet1K 数据集上进行预测推理

我们可以直接加载官方预训练模型在 ImageNet1K 数据集上进行预测推理，需要准备 imagenet 1k的类别索引 json文件，这里我们从github下载即可：
https://github.com/raghakot/keras-vis/blob/master/resources/imagenet_class_index.json

然后准备好部分的 imagenet1K 数据集作为输入的预测图片，最终效果如下图所示：
|775