RTMDet 模型训练
RTMDet (Real-time Models for Object Detection) 是一个高精度、低延时的单阶段目标检测器算法,RTMDet 模型整体结构和 YOLOX 几乎一致,由 CSPNeXt + CSPNeXtPAFPN + 共享卷积权重但分别计算 BN 的 SepBNHead 构成。内部核心模块也是 CSPLayer,但对其中的 Basic Block 改进为了 CSPNeXt Block。
数据集准备
在进行 RTMDet 模型训练之前,我们需要准备好数据集。这里我们以已经标注好的口罩 COCO 数据集为例,您可以在 SSCMA - 公共数据集 中下载该数据集。
模型选择与训练
SSCMA 提供了多种不同的 RTMDet 模型配置,您可以根据自己的需求选择合适的模型进行训练。
rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py
rtmdet_nano_8xb32_300e_coco_relu.py
rtmdet_nano_8xb32_300e_coco_relu_q.py
在此我们以 rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py
为例,展示如何使用 SSCMA 进行 RTMDet 模型训练。
python3 tools/train.py \
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py \
--cfg-options \
data_root=$(pwd)/datasets/coco_mask/mask/ \
num_classes=2 \
train_ann_file=train/_annotations.coco.json \
val_ann_file=valid/_annotations.coco.json \
train_img_prefix=train/ \
val_img_prefix=valid/ \
epochs=150 \
imgsz='(192,192)'
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py
: 指定配置文件,定义模型和训练设置。--cfg-options
: 用于指定额外的配置选项。data_root
: 设定数据集的根目录。num_classes
: 指定模型需要识别的类别数量。train_ann_file
: 指定训练数据的注释文件路径。val_ann_file
: 指定验证数据的注释文件路径。train_img_prefix
: 指定训练图像的前缀路径。val_img_prefix
: 指定验证图像的前缀路径。epochs
: 设置训练的最大周期数。imgsz
:指定模型训练使用的图像尺寸。
等待训练结束后,您可以在 work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6
目录下找到训练好的模型,在查找模型前,我们建议先关注训练结果。以下是对结果的分析以及一些改进方向。
Details
12/19 03:35:57 - mmengine - INFO - Epoch(train) [150][30/30] base_lr: 2.5000e-05 lr: 2.5000e-05 eta: 0:00:00 time: 0.1145 data_time: 0.0051 memory: 383 loss: 0.6947 loss_cls: 0.3424 loss_bbox: 0.3523
12/19 03:35:57 - mmengine - INFO - Saving checkpoint at 150 epochs
12/19 03:35:58 - mmengine - INFO - Evaluating bbox...
Loading and preparing results...
DONE (t=0.30s)
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=0.31s).
Accumulating evaluation results...
DONE (t=0.06s).
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.409
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.930
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.253
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.413
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.437
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.566
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.574
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.579
12/19 03:35:59 - mmengine - INFO - bbox_mAP_copypaste: 0.409 0.930 0.253 -1.000 0.000 0.413
12/19 03:35:59 - mmengine - INFO - Epoch(val) [150][6/6] coco/bbox_mAP: 0.4090 coco/bbox_mAP_50: 0.9300 coco/bbox_mAP_75: 0.2530 coco/bbox_mAP_s: -1.0000 coco/bbox_mAP_m: 0.0000 coco/bbox_mAP_l: 0.4130 data_time: 0.0629 time: 0.1037
通过分析 COCO Eval 的结果可以发现问题所在,并采取相应的措施进行优化,优化方向建议优先从数据集入手,再是训练参数、以及模型结构。
平均精度(AP)方面:
- 在 IoU=0.50:0.95 且 area=all 时,AP 为 0.506,整体处于中等偏低水平,模型在不同交并比综合情况下的检测精度有一定提升空间。
- 当 IoU=0.50 时 AP 达到 0.946,表明在较宽松交并比要求下模型能有较好表现,但 IoU=0.75 时 AP 仅 0.456,意味着模型在高交并比情况下精度很差,对预测框和真实框重合度要求高时表现不佳。
- 按检测目标面积分类来看,area=small 时 AP 为 -1.000,数据异常,小目标检测存在严重问题,数据集的验证集缺乏小目标。而 area=medium 时 AR 和 AP 为 0,表明验证集中包含中等目标的目标,但存在一些其他问题,如训练集中缺乏中等目标、数据增强参数异常等。
平均召回率(AR)方面:
- 在 IoU=0.50:0.95 且 area=all 的不同 maxDets 下,随着 maxDets 从 1 增加到 100,AR 从 0.547 提升到 0.608,增加可检测目标数量上限能在一定程度上提高召回率,但整体数值不算高,模型在实际情况下可能遗漏较多目标。
- 按面积分类中,area=small 的 AR 为 -1.000,再次体现数据集的验证集缺乏小目标的问题。
根据以上数据,我们首先检查数据集中是否由足够的小目标物体、小目标的数据标注是否准确、完整,必要时重新进行标注,确保标注框贴合小目标实际边界,再检查数据集通过训练管线后,经过数据增强后的图像色彩以及标注是否正确、合理。
此外,我们还需要检查训练过程,模型是否收敛等,您可以使用 Tensorboard 来进行查看。
安装并运行时 Tensorboard:
python3 -m pip install tensorboard && \
tensorboard --logdir workdir
在 Scalars 选项卡下,可以查看记录的标量指标(如损失、准确率)随时间(通常是训练轮数)的变化情况。通过观察损失函数的下降趋势和准确率的上升趋势,判断模型是否在正常收敛。如果损失函数不再下降或者准确率不再上升,可能表示模型已经收敛或者出现了问题,这里仅简单介绍一下调整策略。
- 学习率:如果损失函数下降过于缓慢,可以尝试增大学习率;如果损失函数出现剧烈震荡或者不收敛,可能是学习率过大,需要减小学习率。对于学习率的调整策略等请参考 SSCMA - 自定义 - 基础配置结构
- 迭代次数:如果模型在训练过程中还没有完全收敛(例如损失函数仍在下降,准确率仍在上升),可以适当增加迭代次数。如果模型已经收敛,继续增加迭代次数可能会导致过拟合,此时可以减少迭代次数。
在 work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6
目录下找到训练好的模型。此外,当模型训练结果精度不佳时,通过分析 COCO Eval 的结果可以发现问题所在,并采取相应的措施进行优化。
TIP
当模型训练结果精度不佳时,通过分析 COCO Eval 的结果可以发现问题所在,并采取相应的措施进行优化。
模型导出及验证
在训练过程中,您可以随时查看训练日志、导出模型并验证模型的性能,部分模型验证中输出的指标在训练过程中也会显示,因此在这一部分我们会先介绍如何导出模型,然后阐述如何验证导出后模型的精度。
导出模型
这里我们以导出 TFLite 模型为例,您可以使用以下命令导出不同精度的 TFLite 模型:
python3 tools/export.py \
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py \
work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6/epoch_150.pth \
--cfg-options \
data_root=$(pwd)/datasets/coco_mask/mask/ \
num_classes=2 \
train_ann_file=train/_annotations.coco.json \
val_ann_file=valid/_annotations.coco.json \
train_img_prefix=train/ \
val_img_prefix=valid/ \
--imgsz 192 192 \
--format tflite \
--image_path $(pwd)/datasets/coco_mask/mask/valid
WARNING
我们建议在训练和导出时使用相同的分辨率,在当前情况下,使用不同的分辨率训练和导出时,可能导致模型精度降低或完全丢失。
TIP
在导出过程中,可能需要网络环境以安装某些依赖,如果无法访问到互联网,请确保当前的 Python 环境中已经安装了以下依赖:
tensorflow
hailo_sdk_client
onnx
onnx2tf
tf-keras
onnx-graphsurgeon
sng4onnx
onnxsim
此外,onnx2tf
在运行时可能还需要下载 calibration 相关的数据,您可以参考以下链接将其提前下载到 SSCMA 的根目录。
wget https://github.com/PINTO0309/onnx2tf/releases/download/1.20.4/calibration_image_sample_data_20x128x128x3_float32.npy \
-O calibration_image_sample_data_20x128x128x3_float32.npy
验证模型
导出完成后,您可以使用以下命令对 TFLite Int8 模型进行验证:
python3 tools/test.py \
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py \
work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6/epoch_150_int8.tflite \
--cfg-options \
data_root=$(pwd)/datasets/coco_mask/mask/ \
num_classes=2 \
train_ann_file=train/_annotations.coco.json \
val_ann_file=valid/_annotations.coco.json \
train_img_prefix=train/ \
val_img_prefix=valid/ \
imgsz='(192,192)'
得到以下输出:
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.409
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.934
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.284
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.413
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.433
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.559
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.563
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.568
TIP
关于以上输出的详细解释,请参考 COCO 数据集评估指标,在这里我们主要关注 50-95 IoU 和 50 IoU 的 mAP。
QAT
QAT(量化感知训练)是一种在模型训练过程中模拟量化操作,让模型逐步适应量化误差,从而在量化后仍能保持较高精度的方法。SSCMA 支持 QAT,您可以参考以下方法得到 QAT 的模型,并再次验证。
python3 tools/quantization.py \
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py \
work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6/epoch_150.pth \
--cfg-options \
data_root=$(pwd)/datasets/coco_mask/mask/ \
num_classes=2 \
train_ann_file=train/_annotations.coco.json \
val_ann_file=valid/_annotations.coco.json \
train_img_prefix=train/ \
val_img_prefix=valid/ \
imgsz='(192,192)' \
epochs=5
Details
QAT 训练结果:
12/17 09:43:41 - mmengine - INFO - Saving checkpoint at 5 epochs
12/17 09:43:43 - mmengine - INFO - Evaluating bbox...
Loading and preparing results...
DONE (t=0.02s)
creating index...
index created!
Running per image evaluation...
Evaluate annotation type *bbox*
DONE (t=0.31s).
Accumulating evaluation results...
DONE (t=0.06s).
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.600
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.971
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.784
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.605
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.638
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.663
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.663
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.668
12/17 09:43:44 - mmengine - INFO - bbox_mAP_copypaste: 0.600 0.971 0.784 -1.000 0.000 0.605
12/17 09:43:44 - mmengine - INFO - Epoch(val) [5][6/6] coco/bbox_mAP: 0.6000 coco/bbox_mAP_50: 0.9710 coco/bbox_mAP_75: 0.7840 coco/bbox_mAP_s: -1.0000 coco/bbox_mAP_m: 0.0000 coco/bbox_mAP_l: 0.6050 data_time: 0.0342 time: 0.2558
QAT 训练完毕后,会自动导出量化后的模型,您可以使用以下命令对其进行验证:
python3 tools/test.py \
configs/rtmdet/rtmdet_nano_8xb32_300e_coco_ncadc_relu6.py \
work_dirs/rtmdet_nano_8xb32_300e_coco_ncadc_relu6/qat/qat_model_int8.tflite \
--cfg-options \
data_root=$(pwd)/datasets/coco_mask/mask/ \
num_classes=2 \
train_ann_file=train/_annotations.coco.json \
val_ann_file=valid/_annotations.coco.json \
train_img_prefix=train/ \
val_img_prefix=valid/ \
imgsz='(192,192)'
评估得到的精度如下:
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.455
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.842
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.417
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.459
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.499
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.583
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.618
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = -1.000
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.000
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.623