【深度学习】分类卷积神经网络·参数分离的自定义数据集训练

地址：https://gitee.com/MrDreamQ/custom-classification

本人的自制训练集的分类网络训练库——配有轻量型网络训练、部署的详细过程。

网上的不少分类网络的训练库，但许多都没有配置教程，超参数、路径等参数的读取、数据集准备等。因此本人基于MobilenetV2V3的训练库，加入自定义参数及数据集的设置和读取，通过修改配置文件尽可能减少对源文件的修改。

本人自制数据集及损失、准确率示意图如下

本人先使用MobilenetV3_small进行训练，损失值一直在0.1左右，准确率也只有86左右，换至MobilenetV3_large后，损失值下降值0.03，准确率有98左右。

库结构

./Classification/
├── config.yaml 超参数、路径参数
├── data 数据集、生成文件、预训练学习模型存储处
│ ├── dataset 数据集存放 (数据集文件名顺序即为对应标签号)
│ │ ├── runs tensorboard缓存、训练模型保存处
│ │ ├── train
│ │ └── val
│ ├── mobilenet_v3_large.pth mobilenetV3_large 预训练模型
│ ├── mobilenetv3-small.pth mobilenetV3_small 预训练模型
│ └── model 目标模型存储处
│ ├── model.onnx
│ └── network-epoch74_loss0.0347_acc0.986.pth
├── MyDataLoader 自定义数据集加载器
│ ├── MyDataLoader.py
├── network 各类网络模型(增加中…)
│ ├── model_v3.py
├── predict.py 预测文件
├── README.md
├── requirements.txt
├── train.py 训练文件
└── utils 实用工具类、方法存储
├── classDictory.py
├── ConfigReader.py
├── init.py
├── pth2onnx.py
└── shuffle.py

环境准备

在相应conda环境下

pip install -r requirements.txt

需先安装相应的cuda、cudnn环境

数据集准备

将各类别的数据集图像放入各个命名的文件夹中，将各类别文件夹放入data/dataset目录下

由于在自定义的数据集加载器MyDataLoader中使用ImageFolder类，因此每张图像读取后都会附带对应的、以文件夹名为顺序的label

如dataset目录下有

class0

class1

class2

…

则dataset加载后class0的图像label为0, class1的图像label为1, class2的图像label为2…

此方法也免去了一些训练文件中使用文本文件读取图像并手动加标签的麻烦

放置好数据集后，终端或在IDE中运行utils/shuffle.py文件以自动在dataset目录下生成训练集目录train与测试集目录test

# 项目根目录下
python3 utils/shuffle.py

运行后目录中有打乱分离后的各类别数据集的文件夹*(train:val=7:3)*

data目录下可存放多个datasets，想要训练哪个数据集，记得在config.yaml修改dataset_path、save_path即可

训练

参数的读取，数据集的载入、训练、验证等已写好，在config.yaml中配置好参数和路径，项目根目录下运行train.py即可。

注意，文件中的网络模型是mobilenetV3，若有其他网络模型，可以放入network目录中，然后在train.py下做相应修改(包括模型的设置、数据的变换等)，本人之后可能持续更新一些模型，之后将模型选择写入config.yaml中。

训练过程中tensorboard的SummaryWriter缓存会保存在数据集下的runs/目录中

终端下

tensorboard --logdir "data/dataset/runs/"

便能打开tensorboard查看记录的数据。

另外，tensorboard功能很强大，能加入

预测

训练好的模型将保存至dataset/runs/路径下，选择相应模型，移至data/model中存储

然后只需指定相应的pth模型文件路径及图像路径，运行predict.py即可(图像能为数据集的类别文件夹，也可为单张图像)

训练加速相关

根据网上的一些资料，本人在训练库中加入了一些训练加速的措施

如DataLoader设置子线程数num_workers、固定显存pin_memory、cudnn.benchmark、cuda的AMP(自动混合精度)

scaler = torch.cuda.amp.GradScaler()  # 自动混合精度(AMP)
for epoch in range(epochs):
    # 将运算转换为混合精度
    with torch.cuda.amp.autocast():
    logits = net(images.to(device))
    train_loss = loss_function(logits, labels.to(device))
    # 缩放损失，并调用backward()创建缩放梯度
    scaler.scale(train_loss).backward()
    # 取消缩放梯度和调用或跳过optimizer.step()
    scaler.step(optimizer)
    # 为下一次迭代更新比例
    scaler.update()

后续尝试加入更多的分类神经网络，并使用C++进行部署，以高速运行网络
或重构一些目标检测神经网络库并配上教程
遇到问题留言，本人随缘解决
未移植过其他数据集测试，若用cpu训练的话，把与cuda、cudnn相关的全注释掉就好了
喜欢的话给我的库点个star~