上一篇博文我们搭建好了超算平台对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的网络结构,看如何加入自己的修改内容。