本教程旨在介绍如何通过使用PaddleSeg提供的 HRNet
预训练模型在自定义数据集上进行训练、评估和可视化。
我们提前准备好了一份眼底医疗分割数据集,包含267张训练图片、76张验证图片、38张测试图片。通过以下命令进行下载:
python dataset/download_optic.py
接着下载对应的预训练模型
python pretrained_model/download_model.py hrnet_w18_bn_cityscapes
关于已有的HRNet预训练模型的列表,请参见模型组合。如果需要使用其他预训练模型,下载该模型并将配置中的BACKBONE、NORM_TYPE等进行替换即可。
接着我们需要确定相关配置,从本教程的角度,配置分为三部分:
- 数据集
- 训练集主目录
- 训练集文件列表
- 测试集文件列表
- 评估集文件列表
- 预训练模型
- 预训练模型名称
- 预训练模型各阶段通道数设置
- 预训练模型的Normalization类型
- 预训练模型路径
- 其他
- 学习率
- Batch大小
- ...
在三者中,预训练模型的配置尤为重要,如果模型配置错误,会导致预训练的参数没有加载,进而影响收敛速度。预训练模型相关的配置如第二步所展示。
数据集的配置和数据路径有关,在本教程中,数据存放在dataset/optic_disc_seg
中
其他配置则根据数据集和机器环境的情况进行调节,最终我们保存一个如下内容的yaml配置文件,存放路径为configs/hrnet_optic.yaml
# 数据集配置
DATASET:
DATA_DIR: "./dataset/optic_disc_seg/"
NUM_CLASSES: 2
TEST_FILE_LIST: "./dataset/optic_disc_seg/test_list.txt"
TRAIN_FILE_LIST: "./dataset/optic_disc_seg/train_list.txt"
VAL_FILE_LIST: "./dataset/optic_disc_seg/val_list.txt"
VIS_FILE_LIST: "./dataset/optic_disc_seg/test_list.txt"
# 预训练模型配置
MODEL:
MODEL_NAME: "hrnet"
DEFAULT_NORM_TYPE: "bn"
HRNET:
STAGE2:
NUM_CHANNELS: [18, 36]
STAGE3:
NUM_CHANNELS: [18, 36, 72]
STAGE4:
NUM_CHANNELS: [18, 36, 72, 144]
# 其他配置
TRAIN_CROP_SIZE: (512, 512)
EVAL_CROP_SIZE: (512, 512)
AUG:
AUG_METHOD: "unpadding"
FIX_RESIZE_SIZE: (512, 512)
BATCH_SIZE: 4
TRAIN:
PRETRAINED_MODEL_DIR: "./pretrained_model/hrnet_w18_bn_cityscapes/"
MODEL_SAVE_DIR: "./saved_model/hrnet_optic/"
SNAPSHOT_EPOCH: 5
TEST:
TEST_MODEL: "./saved_model/hrnet_optic/final"
SOLVER:
NUM_EPOCHS: 10
LR: 0.001
LR_POLICY: "poly"
OPTIMIZER: "adam"
在开始训练和评估之前,我们还需要对配置和数据进行一次校验,确保数据和配置是正确的。使用下述命令启动校验流程
python pdseg/check.py --cfg ./configs/hrnet_optic.yaml
校验通过后,使用下述命令启动训练
# 指定GPU卡号(以0号卡为例)
export CUDA_VISIBLE_DEVICES=0
# 训练
python pdseg/train.py --use_gpu --cfg ./configs/hrnet_optic.yaml
模型训练完成,使用下述命令启动评估
python pdseg/eval.py --use_gpu --cfg ./configs/hrnet_optic.yaml
使用下述命令启动预测和可视化
python pdseg/vis.py --use_gpu --cfg ./configs/hrnet_optic.yaml
预测结果将保存在visual目录下,以下展示其中1张图片的预测效果:
预训练模型名称 | Backbone | 数据集 | 配置 |
---|---|---|---|
hrnet_w18_bn_cityscapes | HRNet | Cityscapes | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [18, 36] MODEL.HRNET.STAGE3.NUM_CHANNELS: [18, 36, 72] MODEL.HRNET.STAGE4.NUM_CHANNELS: [18, 36, 72, 144] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w18_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [18, 36] MODEL.HRNET.STAGE3.NUM_CHANNELS: [18, 36, 72] MODEL.HRNET.STAGE4.NUM_CHANNELS: [18, 36, 72, 144] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w30_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [30, 60] MODEL.HRNET.STAGE3.NUM_CHANNELS: [30, 60, 120] MODEL.HRNET.STAGE4.NUM_CHANNELS: [30, 60, 120, 240] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w32_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [32, 64] MODEL.HRNET.STAGE3.NUM_CHANNELS: [32, 64, 128] MODEL.HRNET.STAGE4.NUM_CHANNELS: [32, 64, 128, 256] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w40_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [40, 80] MODEL.HRNET.STAGE3.NUM_CHANNELS: [40, 80, 160] MODEL.HRNET.STAGE4.NUM_CHANNELS: [40, 80, 160, 320] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w44_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [44, 88] MODEL.HRNET.STAGE3.NUM_CHANNELS: [44, 88, 176] MODEL.HRNET.STAGE4.NUM_CHANNELS: [44, 88, 176, 352] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w48_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [48, 96] MODEL.HRNET.STAGE3.NUM_CHANNELS: [48, 96, 192] MODEL.HRNET.STAGE4.NUM_CHANNELS: [48, 96, 192, 384] MODEL.DEFAULT_NORM_TYPE: bn |
hrnet_w64_bn_imagenet | HRNet | ImageNet | MODEL.MODEL_NAME: hrnet MODEL.HRNET.STAGE2.NUM_CHANNELS: [64, 128] MODEL.HRNET.STAGE3.NUM_CHANNELS: [64, 128, 256] MODEL.HRNET.STAGE4.NUM_CHANNELS: [64, 128, 256, 512] MODEL.DEFAULT_NORM_TYPE: bn |