在标准数据集上训练预定义模型¶
MMDetection 还提供开箱即用的工具来训练检测模型。本节将介绍如何在标准数据集（例如 COCO）上训练预定义模型（在 configs 中）。
准备数据集¶
准备数据集对于训练也是必要的。有关详细信息，请参见上面的准备数据集部分。
注意：目前，configs/cityscapes下的配置文件使用 COCO 预训练权重进行初始化。如果您的网络连接速度慢或不可用，建议您在开始训练之前下载现有模型，以避免错误。
学习率自动缩放¶
重要：配置文件中的默认学习率适用于 8 个 GPU 和每个 GPU 2 个样本（批大小 = 8 * 2 = 16）。并且它已在config/_base_/schedules/schedule_1x.py中设置为auto_scale_lr.base_batch_size。学习率将根据批大小为 16 的值自动缩放。同时，为了避免影响使用 mmdet 的其他代码库，auto_scale_lr.enable标志的默认设置为False。
如果您想启用此功能，则需要添加参数--auto-scale-lr。并且您需要在处理命令之前检查要使用的配置文件名称，因为配置文件名称指示默认批大小。默认情况下，它是8 x 2 = 16 batch size，例如faster_rcnn_r50_caffe_fpn_90k_coco.py或pisa_faster_rcnn_x101_32x4d_fpn_1x_coco.py。在其他情况下，您会看到配置文件名称中包含_NxM_，例如cornernet_hourglass104_mstest_32x3_210e_coco.py，批大小为32 x 3 = 96，或scnet_x101_64x4d_fpn_8x1_20e_coco.py，批大小为8 x 1 = 8。
请记住检查要使用的特定配置文件的底部，如果批大小不是16，它将包含auto_scale_lr.base_batch_size。如果找不到这些值，请检查位于_base_=[xxx]中的配置文件，您将在其中找到它们。如果您想自动缩放 LR，请不要修改这些值。
学习率自动缩放的基本用法如下。
python tools/train.py \
${CONFIG_FILE} \
--auto-scale-lr \
[optional arguments]
如果您启用了此功能，学习率将根据机器上的 GPU 数量和训练的批大小自动缩放。有关详细信息，请参见线性缩放规则。例如，如果有 4 个 GPU，每个 GPU 上有 2 张图片，lr = 0.01，那么如果有 16 个 GPU，每个 GPU 上有 4 张图片，它将自动缩放为lr = 0.08。
如果您不想使用它，您需要根据线性缩放规则手动计算学习率，然后更改特定配置文件中的optimizer.lr。
在单个 GPU 上训练¶
我们提供tools/train.py来在单个 GPU 上启动训练作业。基本用法如下。
python tools/train.py \
${CONFIG_FILE} \
[optional arguments]
在训练期间，日志文件和检查点将保存到工作目录中，工作目录由配置文件中的work_dir或通过 CLI 参数--work-dir指定。
默认情况下，模型在每个 epoch 在验证集上进行评估，评估间隔可以在配置文件中指定，如下所示。
# evaluate the model every 12 epochs.
train_cfg = dict(val_interval=12)
此工具接受几个可选参数，包括
--work-dir ${WORK_DIR}：覆盖工作目录。--resume：自动从工作目录中的最新检查点恢复。--resume ${CHECKPOINT_FILE}：从特定检查点恢复。--cfg-options 'Key=value'：覆盖已使用配置文件中的其他设置。
注意
resume和load-from之间存在差异
resume加载模型的权重和优化器的状态，并从指定的检查点继承迭代次数，因此训练不会从头开始。另一方面，load-from仅加载模型的权重，并且其训练从头开始。它通常用于微调模型。 load-from需要写入配置文件，而resume作为命令行参数传递。
在 CPU 上训练¶
在 CPU 上训练的过程与单个 GPU 训练一致。我们只需要在训练过程之前禁用 GPU。
export CUDA_VISIBLE_DEVICES=-1
然后运行上面的脚本。
注意:
我们不建议用户使用 CPU 进行训练，因为它太慢了。我们支持此功能以允许用户在没有 GPU 的机器上进行调试，以方便使用。
在多个 GPU 上训练¶
我们提供tools/dist_train.sh来在多个 GPU 上启动训练。基本用法如下。
bash ./tools/dist_train.sh \
${CONFIG_FILE} \
${GPU_NUM} \
[optional arguments]
可选参数与上面说明的一样。
同时启动多个作业¶
如果您想在一台机器上启动多个作业，例如在具有 8 个 GPU 的机器上启动 2 个 4 个 GPU 训练作业，您需要为每个作业指定不同的端口（默认情况下为 29500），以避免通信冲突。
如果您使用dist_train.sh启动训练作业，您可以在命令中设置端口。
CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 ./tools/dist_train.sh ${CONFIG_FILE} 4
CUDA_VISIBLE_DEVICES=4,5,6,7 PORT=29501 ./tools/dist_train.sh ${CONFIG_FILE} 4
使用多台机器训练¶
如果您使用多台机器进行训练，并且这些机器之间仅通过以太网连接，您可以简单地运行以下命令。
在第一台机器上
NNODES=2 NODE_RANK=0 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS
在第二台机器上
NNODES=2 NODE_RANK=1 PORT=$MASTER_PORT MASTER_ADDR=$MASTER_ADDR sh tools/dist_train.sh $CONFIG $GPUS
通常，如果您没有高速网络连接（例如 InfiniBand），速度会很慢。
使用 Slurm 管理作业¶
Slurm 是一个适用于计算集群的良好作业调度系统。在由 Slurm 管理的集群上，您可以使用 slurm_train.sh 来启动训练作业。它支持单节点和多节点训练。
基本用法如下。
[GPUS=${GPUS}] ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR}
以下是用 16 个 GPU 在名为 dev 的 Slurm 分区上训练 Mask R-CNN 的示例，并将工作目录设置为某些共享文件系统。
GPUS=16 ./tools/slurm_train.sh dev mask_r50_1x configs/mask-rcnn_r50_fpn_1x_coco.py /nfs/xxxx/mask_rcnn_r50_fpn_1x
您可以查看 源代码 以查看完整参数和环境变量。
使用 Slurm 时，需要通过以下方式之一设置端口选项
通过
--options设置端口。这种方式更推荐，因为它不会更改原始配置。CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR} --cfg-options 'dist_params.port=29500' CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR} --cfg-options 'dist_params.port=29501'
修改配置文件以设置不同的通信端口。
在
config1.py中，设置dist_params = dict(backend='nccl', port=29500)
在
config2.py中，设置dist_params = dict(backend='nccl', port=29501)
然后您可以使用
config1.py和config2.py启动两个作业。CUDA_VISIBLE_DEVICES=0,1,2,3 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config1.py ${WORK_DIR} CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} config2.py ${WORK_DIR}
使用自定义数据集训练¶
在本部分，您将了解如何使用自定义数据集训练预定义模型，然后对其进行测试。我们以 气球数据集 为例来描述整个过程。
基本步骤如下
准备自定义数据集
准备配置
在自定义数据集上训练、测试和推理模型。
准备自定义数据集¶
有三种方法可以在 MMDetection 中支持新的数据集
将数据集重新组织为 COCO 格式。
将数据集重新组织为中间格式。
实现一个新的数据集。
通常，我们建议使用前两种方法，它们通常比第三种方法更容易。
在本说明中，我们提供了一个将数据转换为 COCO 格式的示例。
注意：自 MMDetection 3.0 以来，数据集和指标已解耦，CityScapes 除外。因此，用户可以在验证期间对任何格式的数据集使用任何类型的评估指标。例如：使用 VOC 指标评估 COCO 数据集，或使用 VOC 和 COCO 指标评估 OpenImages 数据集。
COCO 注释格式¶
实例分割的 COCO 格式所需键如下所示，有关完整详细信息，请参阅 此处。
{
"images": [image],
"annotations": [annotation],
"categories": [category]
}
image = {
"id": int,
"width": int,
"height": int,
"file_name": str,
}
annotation = {
"id": int,
"image_id": int,
"category_id": int,
"segmentation": RLE or [polygon],
"area": float,
"bbox": [x,y,width,height], # (x, y) are the coordinates of the upper left corner of the bbox
"iscrowd": 0 or 1,
}
categories = [{
"id": int,
"name": str,
"supercategory": str,
}]
假设我们使用气球数据集。下载数据后，我们需要实现一个函数将注释格式转换为 COCO 格式。然后我们可以使用实现的 CocoDataset 加载数据并执行训练和评估。
如果您看一下数据集，您会发现数据集格式如下所示
{'base64_img_data': '',
'file_attributes': {},
'filename': '34020010494_e5cb88e1c4_k.jpg',
'fileref': '',
'regions': {'0': {'region_attributes': {},
'shape_attributes': {'all_points_x': [1020,
1000,
994,
1003,
1023,
1050,
1089,
1134,
1190,
1265,
1321,
1361,
1403,
1428,
1442,
1445,
1441,
1427,
1400,
1361,
1316,
1269,
1228,
1198,
1207,
1210,
1190,
1177,
1172,
1174,
1170,
1153,
1127,
1104,
1061,
1032,
1020],
'all_points_y': [963,
899,
841,
787,
738,
700,
663,
638,
621,
619,
643,
672,
720,
765,
800,
860,
896,
942,
990,
1035,
1079,
1112,
1129,
1134,
1144,
1153,
1166,
1166,
1150,
1136,
1129,
1122,
1112,
1084,
1037,
989,
963],
'name': 'polygon'}}},
'size': 1115004}
注释是一个 JSON 文件，其中每个键表示图像的所有注释。将气球数据集转换为 coco 格式的代码如下所示。
import os.path as osp
import mmcv
from mmengine.fileio import dump, load
from mmengine.utils import track_iter_progress
def convert_balloon_to_coco(ann_file, out_file, image_prefix):
data_infos = load(ann_file)
annotations = []
images = []
obj_count = 0
for idx, v in enumerate(track_iter_progress(data_infos.values())):
filename = v['filename']
img_path = osp.join(image_prefix, filename)
height, width = mmcv.imread(img_path).shape[:2]
images.append(
dict(id=idx, file_name=filename, height=height, width=width))
for _, obj in v['regions'].items():
assert not obj['region_attributes']
obj = obj['shape_attributes']
px = obj['all_points_x']
py = obj['all_points_y']
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
poly = [p for x in poly for p in x]
x_min, y_min, x_max, y_max = (min(px), min(py), max(px), max(py))
data_anno = dict(
image_id=idx,
id=obj_count,
category_id=0,
bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
area=(x_max - x_min) * (y_max - y_min),
segmentation=[poly],
iscrowd=0)
annotations.append(data_anno)
obj_count += 1
coco_format_json = dict(
images=images,
annotations=annotations,
categories=[{
'id': 0,
'name': 'balloon'
}])
dump(coco_format_json, out_file)
if __name__ == '__main__':
convert_balloon_to_coco(ann_file='data/balloon/train/via_region_data.json',
out_file='data/balloon/train/annotation_coco.json',
image_prefix='data/balloon/train')
convert_balloon_to_coco(ann_file='data/balloon/val/via_region_data.json',
out_file='data/balloon/val/annotation_coco.json',
image_prefix='data/balloon/val')
使用上述函数，用户可以成功地将注释文件转换为 json 格式，然后我们可以使用 CocoDataset 使用 CocoMetric 训练和评估模型。
准备配置¶
第二步是准备一个配置，以便可以成功加载数据集。假设我们想使用带有 FPN 的 Mask R-CNN，用于在气球数据集上训练检测器的配置如下所示。假设配置位于目录 configs/balloon/ 下，并命名为 mask-rcnn_r50-caffe_fpn_ms-poly-1x_balloon.py，配置如下所示。请参阅 了解配置 - MMDetection 3.0.0 文档 以获取有关配置文件的详细信息。
# The new config inherits a base config to highlight the necessary modification
_base_ = '../mask_rcnn/mask-rcnn_r50-caffe_fpn_ms-poly-1x_coco.py'
# We also need to change the num_classes in head to match the dataset's annotation
model = dict(
roi_head=dict(
bbox_head=dict(num_classes=1), mask_head=dict(num_classes=1)))
# Modify dataset related settings
data_root = 'data/balloon/'
metainfo = {
'classes': ('balloon', ),
'palette': [
(220, 20, 60),
]
}
train_dataloader = dict(
batch_size=1,
dataset=dict(
data_root=data_root,
metainfo=metainfo,
ann_file='train/annotation_coco.json',
data_prefix=dict(img='train/')))
val_dataloader = dict(
dataset=dict(
data_root=data_root,
metainfo=metainfo,
ann_file='val/annotation_coco.json',
data_prefix=dict(img='val/')))
test_dataloader = val_dataloader
# Modify metric related settings
val_evaluator = dict(ann_file=data_root + 'val/annotation_coco.json')
test_evaluator = val_evaluator
# We can use the pre-trained Mask RCNN model to obtain higher performance
load_from = 'https://download.openmmlab.com/mmdetection/v2.0/mask_rcnn/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth'