【超级视客营】基于超算平台的MMYOLO实践过程记录（自定义数据集实现YOLO v5）

上一篇博文我们搭建好了超算平台对MMYOLO编译环境，并通过单张图片加载训练好的模型进行推理。本篇文章参考MMYOLO官方发布的自定义数据集 标注+训练+测试+部署 全流程文档，进一步在超算平台中实现yolo v5网络的训练。

1.数据集的构建和可视化分析

1.1 MIDOG公开数据集介绍

本次实验采用的是病理图像有丝分裂检测数据集，边界框标注了两个类别： mitosis figure表示该细胞为有丝分裂细胞，not mitosis figure为比较容易与有丝分裂状态混淆的其他阶段的细胞。关于该数据集的详细介绍可以参考之前的两篇文章。

• 有丝分裂检测竞赛，MICCAI举办。

• 200张训练图像，每张图片大小为组织切片感兴趣区域，尺寸大约为5k*5k pix

• 染色标准相同：UMC Utrecht病理实验室制备得到

• 扫描设备不同：50 each from Hamamatsu XR nanozoomer 2.0 (HXR), Hamamatsu S360 0.5 NA (HS360), Aperio ScanScope CS2 (ACS), and Leica GT450

1.2 数据格式的转换

• 解析MIDOG的标注文件.json → df = get_bbox_df()

• 训练集和测试集划分 → train_test_split(df)

• 图像块生成：不重叠裁剪成512*512尺寸的patch → mask_segmentor()

from PatchSeg_functions import mask_segmentor
from Data_functions import get_bbox_df, train_test_split
############ Preparation ################## Get dataset and generate 512x512 crops from the original WSIs
df = get_bbox_df()
df_train, df_test = train_test_split(df)
image_ids=list(df['file_name'].unique())
mask_segmentor(image_ids,df,512,512,categories=['mitotic figure'],
               image_folder="I:/Pathology/Dataset/MIDOG2021/image_crops/")

• 解析json文件

• 主要是两个列表元素：images和annotations

• images列表中包含200张图像信息：file_name，id，width和height

• annotation中包含了4435个标注信息：bbox坐标，category_id，image_id，id（第几个标注）

• 最终得到df，shape = {tuple: 2}(4435, 8)，如下图所示

1.3 训练集和测试集划分

• A、B、C三个中心分别有50张图片，每个中心以8：2的比例划分训练集和测试集【前40个为训练集，后10个为测试集】

• 分别训练三个模型，并在另外1个中心进行外部测试。eg：A训练模型，B和C外部测试

• 只对外部测试集进行校正，参与模型训练的数据集不进行颜色校正

• 划分好的数据集文件夹如下图，和coco格式保持一致

1.4 检查数据集的标签（重要）

数据格式准备好之后，可以使用tools中的分析工具browse_coco_json.py验证一下自己的数据集标签是否有问题。具体命令如下：

python tools/analysis_tools/browse_coco_json.py --img-dir ${图片文件夹路径} \
                                                --ann-file ${COCO label json 路径}

另外MMYOLO还提供了数据集分布可视化的工具，非常好用。该脚本使用方式如下：

python tools/analysis_tools/dataset_analysis.py ${CONFIG} \
                                                [--val-dataset ${TYPE}] \
                                                [--class-name ${CLASS_NAME}] \
                                                [--area-rule ${AREA_RULE}] \
                                                [--func ${FUNC}] \
                                                [--out-dir ${OUT_DIR}]

2.修改config文件

接下来可以为自己的数据集和网络建立一个新的config文件，放置在config文件夹下面。以下是我为该数据集创建的配置文件。主要修改了这些内容：

• data_root数据路径：在超算平台中数据集都放置在data里（这里提供了300G空间，并且放在这个路径下数据读取速度也快），可根据具体位置设置

• load_from预训练模型的下载路径：这里注意不能放置下载链接，因为超算平台中的计算节点是无法链接互联网的，需要把模型下载好，改成存放路径

• class_name类别信息、数据集中的路径信息需要修改成自己的

• anchors尺寸需要修改，MMYOLO也提供了很多生成anchor的脚本，这里采用v5-k-means聚类的方式进行锚框的优化。命令如下：

python tools/analysis_tools/optimize_anchors.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py \
                                                --algorithm v5-k-means \
                                                --input-shape 640 640 \
                                                --prior-match-thr 4.0 \
                                                --out-dir work_dirs/dataset_analysis_cat

3.编写sh脚本启动训练

3.1 运行前注意事项

修改好配置文件之后，就可以使用tools中的train，因为超算平台无法在登录节点运行含有GPU内容的脚本，所以train.py无法在登录节点运行

• 我尝试运行过，它会提示mmengine库不存在等问题，实际上已经安装了，只有在计算节点才可以识别配置的第三方库。

3.2 sh脚本

• 首先，我把.sh文件写在了与mmyolo文件夹并列的目录中，具体内容如下：

#!/bin/sh
# 加载 anaconda
module load anaconda/2020.11
# 加载 cuda 11.3
module load cuda/11.3
# 激活 python 虚拟环境
source activate openmmlab
cd mmyolo
export PYTHONUNBUFFERED=1
# yolov5 python 训练的的命令
python tools/train.py configs/custom_dataset/yolov5_s-v61_syncbn_fast_1xb32-100e_cat.py 

• 接下来在登录节点输入命令： sbatch --gpus=1 sub.sh启动训练过程，这时候本应该在终端显示我们的配置文件、编译环境、网络结构、训练epoch等内容，超算平台将这些内容全部输出到slurm-276560.out文件中，我们可以通过查看该文件的内容对代码进行调试。

• 至此，网络就正常跑通了，接下来等训练结束，可以进行测试。后续开始学习yolo v5的网络结构，看如何加入自己的修改内容。