Pytorch复现RepVGG模型，实现花分类

写在开头，本篇博客是博主跟着一位b站up主学习过程中，发现up主只对网络讲解，但没有复现的视频。
在此放上up主的链接(点我)

为了加深印象，复现是不可少的，所有博主就去论文源码扒代码去了，但源码其实很多东西不需要的，就对于我要实现花分类任务而言，所以接下来所放的代码，博主都经过了一定的修改，也会写一些改代码时候报错的心得
论文：《RepVGG: Making VGG-style ConvNets Great Again》
up讲解视频:RepVGG网络讲解
源码：这里给的镜像，因为github有时候进不去

一、环境安装和项目克隆

1.1 环境安装

• Anaconda3（建议使用）
• python=3.6/3.7/3.8
• pycharm (IDE)
• pytorch=1.11.0 (pip package)
• torchvision=0.12.0 (pip package)
• cudatoolkit=11.3(pip下载pytorch时)
• Wandb

注意： wandb安装教程，应该到输入密钥那就可以了。

如果有其他库报错，显示没有这个库，自己pip安装一下
如果不会安装环境这里再推荐一篇博客(点我)，里面讲的很详细了，跟着安装就行了

注意：这里再提一嘴，不一定花分类，只要是图像分类任务，都可以用此代码，

插一块验证pytorch是否安装好的代码，自己创建一个新文件复制进去运行

import torch
a = torch.cuda.is_available()
print(a)
ngpu= 1
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())

运行后的结果

1.2 项目克隆

这里放上博主实现的代码，百度云链接
链接：https://pan.baidu.com/s/1-O5HMkHpa7U15Hb5SgIqqw?pwd=lgr9
提取码：lgr9

1.3 项目整体介绍

将下载的压缩包解压，用pycharm打开

这里简单的介绍下各个文件的功能，这里只讲个大概，知道干嘛的就行了，因为博主也不可能对代码全部了解

• 最简单的，3张图片用来给预测脚本用的(1.jpg，2.jpg，R.jpg)
• save_model:用来存训练后和转换后的权重
• wandb：train.py使用了wandb库，用来记录模型训练过程，这个文件库自动产生保存一些日志文件的
• class_indices.json:生成类别名称以及对应的数字索引
• convert.py:用于权重转化
• model.py:模型搭建的文件
• my_dataset.py:用于处理数据集的函数
• predict.py:预测脚本
• Re-parameterization.py：用于验证论文的想法是否正确，可以不管
• se_block.py:保存SE模块(Squeeze-and-Excite)，不懂就百度搜下
• train.py:训练脚本
• utils.py:工具函数，存放一些有用的函数，比如训练和验证的函数

以上就是博主项目的整体介绍

二、数据集和预训练

2.1 数据集的准备

博主复现图像分类任务的数据集一般都用的5类花分类
你要用自己的分类数据集也可以用这个代码，很简单，按图中分类的格式放就行，文件名字就是类别名字，至于LICENSE.txt文件不用管
花分类数据集链接
下载后解压里面应该是这样的，5类花，分别是雏菊、蒲公英、玫瑰、太阳花、郁金香

下载后放着就行了，啥也不要弄
下面是图像总数量，按(8:2)划分的训练集和验证集数量，后面会讲如何修改比例

2.2 获得预训练权重

使用预训练权重目的是为了缩短模型的训练时间，并且能达到比较好的精度
这里源码给了很多不同配置的RepVGG，博主选择不大不小的RepVGG-A2
这里博主放上RepVGG-A2的预训练权重链接
链接：https://pan.baidu.com/s/1ezjFcjzHAy7TxIrE2KloGg?pwd=tr3q
提取码：tr3q
将权重放到你自己喜欢的位置

三、训练花分类模型

3.1 相关文件配置

验证model.py文件

你可以执行运行model.py文件，博主在里面添加了统计模型参数的函数
拉到最后就可以看到
其他模型也可以用这个统计参数的代码，修改create_RepVGG_A2改成其他模型函数的名字即可

if __name__ == "__main__":
    a = torch.rand(1, 3, 224, 224)
    model = create_RepVGG_A2(num_classes=5) # 花分类是5类这里写5
    y = model(a)
    print(y)
    # 统计模型参数
    sum = 0
    for name, param in model.named_parameters():
        num = 1
        for size in param.shape:
            num *= size
        sum += num
        # print("{:30s}:{}".format(name,param.shape))
    print("total param num {}".format(sum))  # total param num 26,808,645

你可以验证下是不是跟博主的参数一样，# total param num 26,808,645
这里再给上论文给RepVGG-A2的参数，这里的Top-1的精度是在ImageNet 数据集上的精度

可以看到论文上是25.49M,跟我们代码跑出来的(26.80M)稍微有点出入，但感觉不大

utils.py文件

里面第一函数，val_rate：这个就是验证集的比例(现在是8：2)，你想多划点改这个就行

def read_split_data(root: str, val_rate: float = 0.2):

3.2 准备训练模型

如果需要训练其他配置的模型，找到train.py文件，把下面代码create_RepVGG_A2改了即可，应该是第9行，博主model.py文件就放了2个配置的RepVGG，另一个是RepVGG_B1g2，可以去model文件拉到后面看

from model import create_RepVGG_A2 as create_model

接着拉到最底，找到main函数入口

if __name__ == '__main__':
"""
    opt模型主要参数解析：
    --num_classes：数据集类别数
    --epochs: 训练轮数
    --batch-size: 每次喂入多少图片
    --lr: 学习率
    --wd: 权重衰减
    --data-path: 数据集路径
    --weights: 预训练权重路径，空为不加载预训练权重
    --freeze-layers: 冻结权重，True为只训练最后一层全连接层
    --device: 显卡编号，默认0第一张
    --project：项目名称
    --model_name: 模型配置
    --save: 权重保存路径
"""
    parser = argparse.ArgumentParser(description='RepVGG config')
    parser.add_argument('--num_classes', type=int, default=5)
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch-size', type=int, default=16)
    parser.add_argument('--lr', type=float, default=0.01)  
    parser.add_argument('--wd', type=float, default=1e-5)  
    parser.add_argument('--data-path', type=str,
                        default="../flower_photos") # ..为上一层
    parser.add_argument('--weights', type=str, default="",
                        help='initial weights path')  # 空为不加载预训练权重
    parser.add_argument('--freeze-layers', type=bool, default=False)  # 只训练全连接层
    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
    parser.add_argument('--project', type=str, default='RepVGG')
    parser.add_argument('--model_name', type=str, default='RepVGG-A2')
    parser.add_argument('--save', type=str, default='save_model/RepVGG-A2.pth')
    opt = parser.parse_args()

如果爆显存了把batch-size调低，8，4，2这样调。调最低都没用的话，就是该换显卡了。
其他参数根据自己情况看着调就行了
改好后环境什么都没问题，直接运行就行了
训练时候会显示每轮次训练集和验证集的loss，acc如图所示

注意： 如果是因为wandb报错，而且怎么都不会改的话，就学图片中把wandb注释了，不用了，博主是因为wandb能自动将绘制曲线上传云端，觉得挺方便就用了

这里再给个小技巧，如果还没正式训练，而是在调试的时候，也把wandb的几行代码注释了，这样调试的时候会快点，而且你如果调试不关wandb，会把每次调试的结果也传到云端，到时候还是得手动删也麻烦

3.3 查看训练结果

运行完后可以到save_model文件夹查看保存的权重

如果wandb每报错，训练结束后，终端会有一个链接，点进去就可以看到你这次的训练结果了
这里先展示加载预训练权重，epochs=10，只训练全连接层的结果，模型名字pre-RepVGG-A2

如果你的x坐标是step而不是epoch的话，点击图中的小铅笔，将x选为epoch，这样就跟博主一样了

如果要在同个project下运行两个不同配置的模型，需要改你的model_name,比如我现在第一个是加载了预训练权重的所有叫pre-RepVGG-A2

parser.add_argument('--model_name', type=str, default='pre-RepVGG-A2')
parser.add_argument('--weights', type=str, default="你的预训练权重路径",help='initial weights path')
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--freeze-layers', type=bool, default=True)

下面我要训练一个不加载预训练全重，epoch=30，训练全部权重

parser.add_argument('--model_name', type=str, default='RepVGG-A2')
parser.add_argument('--weights', type=str, default="",help='initial weights path')
parser.add_argument('--epochs', type=int, default=30)
parser.add_argument('--freeze-layers', type=bool, default=False)

这样做你就可以比较了，按左上角

你就能看到两个不同模型的精度比较图了

当然在这个项目其实没有比较的意义，博主只是演示下有这个操作

3.4 转换权重

这里是RepVGG独特的，因为你如果看过论文可以知道，这个模型在训练和推理的时候用的不是同一个模型，附上论文原图，这样做的目的是加快推理速度，当然这就不讲原理了，这已经超过本篇博客的目的了，详情去看下论文就懂了，这里大家只要知道就ok了。

打开文件

这里要改的东西不多，就一个是之前训练后保存的权重路径，一个是保存转换后的权重的路径

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='RepVGG(plus) Conversion')
    parser.add_argument('--load', default='save_model/RepVGG-A2.pth', help='path to the weights file')
    parser.add_argument('--save', default='save_model/RepVGG-A2-deploy.pth', help='path to the weights file')

改好运行会显示，加载转换前的权重，以及转换后权重，对于同一个输入的输出是否一致，如果显示0表示一样，没问题。

并在save_model文件中产生转换后的权重也就是RepVGG-A2-deploy.pth

四、推理测试

打开predict.py文件，要改的也不多，一个是对要测试图片的路径，另一个是转换后权重的路径
这里图片可以你自己去网上找一张，也可以用博主项目里自带的，应该有3张图片可以用于，我这里图片是从验证集里随便找的3张，这个只要不是训练集的图片就行。

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--image_path', type=str, default='R.jpg')
    parser.add_argument('--weights_path', type=str, default='save_model/RepVGG-A2-deploy.pth')

运行脚本，结果如图所示，对于输入图片郁金香的测试准确率有96%,推理速度是3.2s多

五、最后

其实到这里基本结束了，这里我们在验证一下之前说的，论文里训练和推理的模型结构是不一样的，论文说推理的结果速度更快
然后这里需要改一下代码，将depoly改为False

# 39行
model = create_model(num_classes=5, deploy=True).to(device)

然后改权重路径，改为没有经过转换的权重路径，将RepVGG-A2-deploy.pth改为RepVGG-A2.pth

    parser.add_argument('--save', default='save_model/RepVGG-A2.pth', help='path to the weights file')

再次运行代码

可以看到准确率其实差不多，花费时间为3.3s，大概快了3%吧，看来论文说的是真的doge,哈哈。

注意：怕前面有人略过直接看结尾，这个项目里的代码适用其他图像分类任务，也就是是说，你随便换数据集照样管用。
还有train.py,predict.py,my_dataset.py,utils.py这4个文件适用其他的图像分类模型，比如swin transformer,ConvNeXt等等模型，只要是图像分类任务的。当然需要一定的改动，但改动应该不多，因为博主的其他模型都是用的这个代码，具体细节这就不多讲了，靠大家自己了。
model.py文件肯定是不能用的，每个不同的模型，都不一样。

可能有些细节每写全，不过我已经尽可能把我能考虑到的都写到这篇博客了

最后，如果这篇文章对你有帮助请点点赞吧，毕竟改源码也改了一天了，码字也不容易。
最后的最后，加油吧各位。