【图像分类】RepVGG快速实现跑通

官方仓库

RepVGG github仓库地址

克隆指令：

git clone https://github.com/DingXiaoH/RepVGG.git

一、数据集准备

数据集结构：

--- data								# 数据集根目录
	--- train							# 训练集
		--- class1						# 训练集子文件夹名称即为类别
			--- xxx.jpg					# 某类别下的某图片
	---val								# 测试集
		--- class1
			--- xxx.jpg

• train与val下的文件夹个数与名称都需要保持一致
• 子文件夹的名称即为图片类别标签
  如下图所示
  data = “*/TSR”, 每个子文件夹命名为类别名称
  

二、参数修改

trian.py中的参数修改

train.py的参数在整个脚本最上面，找到如下段落

IMAGENET_TRAINSET_SIZE = 1281167		#修改为训练集大小

parser.add_argument('data', ... default= './my_dataset')	# default= 训练集路径
...
parser.add_argument('--epochs', default=120,...,)			# 训练轮数，为了验证是否能够跑通设置为个位数就可以
...
parser.add_argument('-b', '--batch-size', default=256,...') # 根据gpu修改，先默认不修改，如果报错xxx out of memory, 再修改为 128， 依次类推 64、32、16、8、4等

参数修改详细说明：

• IMAGENET_TRAINSET_SIZE
  这个参数表示训练集的图片量，默认值1281167是 ImageNet 数据集 2012 版本中的训练集图像数量，修改为自己的训练集大小即可

• data：即 parser.add_argument(‘data’, metavar=‘DIR’, default= ‘./my_dataset’,help=‘path to dataset’) 代码中的默认值 default = xxx， 将其修改为数据集路径即一中所述data；

• arch: 默认为RepVGG-A0，在repvgg.py中有许多模型结构，先使用默认值；

• –epochs：第一次训练设置default = 5就好；

• –b, --batch-size：值越大训练越快，但所需内存也越大，如果提示out of memory，慢慢减小即可
  其他参数暂时不用修改

注：修改repvgg.py中 create_RepVGG_A0 的 num_classes为数据集类别
（前面arch参数指定了RepVGG_A0 ，故修改对应函数类别数量）

def create_RepVGG_A0(deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=37,
                  width_multiplier=[0.5, 0.75, 0.75, 1.5], override_groups_map=None, deploy=deploy)

我的数据集有37类，故设置num_classes=37

三、训练

在train.py对应的终端中输入以下命令，即可开始训练

python train.py

训练完成后会在train.py同级目录下生成.pth.tar文件，这便是训练好的权重了