RepVGG：使用自己的数据集进行训练、模型转换、推理

关于RepVGG的模型介绍以及论文、代码的地址，可以参考这篇博客。

本篇博客的目的，是介绍如何使用自己的数据集训练一个RepVGG模型。

目录

1、数据集的准备

2、模型的训练

3、模型转换

---

1、数据集的准备

由于是分类模型，所以数据集需要准备至少两个类别，这里假设为：A、B两类。

源码中是假设数据集分为两部分：train、val，分别对应训练、验证集；在train或者val数据集中，每类样本分别存放于一个独立文件夹，也即：A类的样本存放于文件夹A中，B类的样本存放于文件夹B中；因此，总的数据集为：train、val中分别有A、B两个文件夹。

如果你的数据集是按照这种格式存储的，就可以直接使用源码进行训练了，因为源码中的dataset是使用的torch中的自带数据集定义：

import torchvision.datasets as datasets    
train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))

但是，如果你的dataset是自定义的，需要把对应的这部分代码修改一下，改为你重写的dataset类。

例如，我的数据集是将每个图片的地址存放到对应的TXT文件中：

 每个TXT文件中存放的格式为：

D:/dataset/A/VID_20200424_153035_frame_15747_0_0.99.jpg	0
D:/dataset/A/VID_20200424_153035_frame_15776_0_0.99.jpg	0
D:/dataset/B/1000_0_0.54.jpg	1
D:/dataset/B/1001_0_0.51.jpg	1

因此，我需要自定义dataset类，并重写了__len__()、__getitem__()方法，然后将源码训练部分中的dataset改为我自己定义的类。

这里贴出我的自定义dataset类：

import os
import numpy as np
from PIL import Image

from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms


class ClassifierLoder(Dataset):
    def __init__(self,
                 train_val_files_path="train_val_files",
                 mode="train",
                 transform=None
                 ):
        super(ClassifierLoder, self).__init__()
        self.transform = transform
        self._txt_file = open(os.path.join(train_val_files_path, mode + '.txt'), 'r')
        self._lines = self._txt_file.readlines()
        np.random.shuffle(self._lines)
        self.imgs = [line.split('\t')[0] for line in self._lines]
        self.labels = [int(line.split('\t')[-1].replace('\n', '')) for line in self._lines]

    def __len__(self):
        return len(self._lines)

    def __getitem__(self, idx):
        img_asbp = self.imgs[idx]
        label = self.labels[idx]
        img = Image.open(img_asbp)

        if self.transform:
            img = self.transform(img)
        else:
            # Resize和ToTensor是必须的，不然尺寸不统一且不是torch.tensor格式
            img = transforms.Resize([192, 118])(img)
            img = transforms.ToTensor()(img)

        return img, label

对应的，源码train.py中，dataset的定义我就改为：

    train_dataset = ClassifierLoder(train_val_files_path=args.train_dateset_path,
                                         mode="train",
                                         transform=transform_train)

2、模型的训练

数据准备完毕并修改对应dataset对象之后，需要对训练代码微调下。

需要注意的一点是，源码是使用ImageNet数据集进行训练的，有个全局变量IMAGENET_TRAINSET_SIZE代表的是数据集大小，需要改为你的数据集大小，不然影响学习率调度函数的使用，可以将其定义放到DataLoader之后，然后令其等于len(train_loader)即可。

其他的，就是那些超参了，修改一下即可训练。

这一步没啥好说的，轻微修改一下，训练即可。

3、模型转换

由于RepVGG是训练-推理解耦的，也即其训练时的网络结构和推理时的网络结构是不同的：训练阶段，网络包含了残差结构、不同大小的卷积核（3*3、1*1）；而推理阶段，则只包含3*3卷积且为plain结构。

因此，训练之后保存的模型，需要先进行模型转换，才能用于推理。转换代码如下：

from repvgg import repvgg_model_convert, create_RepVGG_A0
train_model = create_RepVGG_A0(deploy=False)
train_model.load_state_dict(torch.load('RepVGG-A0-train.pth'))          # or train from scratch
# do whatever you want with train_model
deploy_model = repvgg_model_convert(train_model, create_RepVGG_A0, save_path='repvgg_deploy.pth')

或者：

deploy_model = create_RepVGG_A0(deploy=True)
deploy_model.load_state_dict(torch.load('RepVGG-A0-deploy.pth'))

上述转换函数的定义在repvgg.py中，其参数包括已训练好的模型、训练阶段的模型定义、推理阶段模型保存路径。

可以先导入训练好的模型，然后通过转换函数得到推理模型，然后进行推理；也可以在转换之后，先定义一个推理阶段的模型，然后导入转换好的推理模型。

其中，关于模型定义，有两种模式：训练模式、推理模型，通过参数deploy进行控制，当为TRUE时，代表推理阶段，为FALSE时，代表训练阶段。