Deformable DETR检测模型训练

这篇教程以Deform-detr为例,告诉大家如何使用HAT算法包训练一个定点的检测模型。 在开始量化感知训练,也就是定点模型训练之前,首先需要训练一个精度较高的纯浮点模型,然后基于这个纯浮点模型做finetune,就可以快速的训练出定点模型。 所以我们从训练一个纯浮点的Deform-detr模型开始讲起。

训练流程

和其他任务一样,对于所有的训练,评测任务,HAT统一采用 tools + config 的形式来完成。在准备好原始数据集之后,可以通过下面的流程,方便地完成整个训练的流程。

数据集准备

在开始训练模型之前,第一步是需要准备好数据集,这里我们下载MSCOCO的 train2017.zipval2017.zip 做为网络的训练集和验证集,同时需要下载相应的标签数据 annotations_trainval2017.zip,解压缩之后数据目录结构如下所示:

tmp_data |-- mscoco -- annotations_trainval2017.zip |-- train2017.zip |-- val2017.zip |-- annotations |-- train2017 |-- val2017

同时,为了提升训练的速度,我们对原始的jpg格式的数据集做了一个打包,将其转换成lmdb格式的数据集。只需要运行下面的脚本,就可以成功实现转换:

python3 tools/datasets/mscoco_packer.py --src-data-dir ./tmp_data/mscoco/ --target-data-dir ./tmp_data/mscoco --split-name train --pack-type lmdb python3 tools/datasets/mscoco_packer.py --src-data-dir ./tmp_data/mscoco/ --target-data-dir ./tmp_data/mscoco --split-name val --pack-type lmdb

上面这两条命令分别对应着转换训练数据集和验证数据集,打包完成之后,data目录下的文件结构应该如下所示:

tmp_data |-- mscoco |-- annotations |-- train2017 |-- train_lmdb |-- val2017 |-- val_lmdb

train_lmdbval_lmdb 就是打包之后的训练数据集和验证数据集,也是网络最终读取的数据集。

模型训练

数据集准备好之后,就可以开始训练浮点型的Deformable Detr检测网络了。在网络训练开始之前,你可以使用以下命令先测试一下网络的计算量和参数数量:

python3 tools/calops.py --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py --input-shape "1,3,800,1332"

如果你只是单纯的想启动这样的训练任务,只需要运行下面的命令就可以:

python3 tools/train.py --stage float --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py python3 tools/train.py --stage calibration --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py python3 tools/train.py --stage qat --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

以上命令分别完成浮点模型和定点模型的训练,其中定点模型的训练需要以训练好的浮点模型为基础,具体内容请阅读 量化感知训练 章节的内容。

导出定点模型

完成量化训练后,便可以开始导出定点模型。可以通过下面命令来导出:

python3 tools/export_hbir.py --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

模型验证

在完成训练之后,可以得到训练完成的浮点、量化或定点模型。 和训练方法类似,我们可以用相同方法来对训好的模型做指标验证,得到为 FloatCalibrationQatQuantized 的指标:

python3 tools/predict.py --stage "float" --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py python3 tools/predict.py --stage "calibration" --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py python3 tools/predict.py --stage "qat" --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

和训练模型时类似,--stage 后面的参数为 "float""calibration""qat" 时,分别可以完成对训练好的浮点模型、量化模型的验证。

定点模型精度验证可使用下面命令,但需要注意是必须要先导出hbir:

python3 tools/predict.py --stage "int_infer" --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

模型推理

HAT 提供了 infer_hbir.py 脚本提供了对定点模型的推理结果进行可视化展示:

python3 tools/infer_hbir.py --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py --model-inputs img:${img-path} --save-path ${save_path}

仿真上板精度验证

除了上述模型验证之外,我们还提供和上板完全一致的精度验证方法,可以通过下面的方式完成:

python3 tools/validation_hbir.py --stage "align_bpu" --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

定点模型检查和编译

在HAT中集成的量化训练工具链主要是为了地平线的计算平台准备的,因此,对于量化模型的检查和编译是必须的。 我们在HAT中提供了模型检查的接口,可以让用户定义好量化模型之后,先检查能否在 BPU 上正常运行:

python3 tools/model_checker.py --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

在模型训练完成后,可以通过 compile_perf_hbir 脚本将量化模型编译成可以上板运行的 hbm 文件,同时该工具也能预估在 BPU 上的运行性能:

python3 tools/compile_perf_hbir.py --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

以上就是从数据准备到生成量化可部署模型的全过程。

训练细节

模型结构

Deformable DETR的网络结构可以参考 论文 ,这里不做详细介绍。我们通过在config配置文件中定义 model 这样的一个dict型变量,就可以方便的实现对模型的定义和修改。

model = dict( type="DeformableDETR", backbone=dict( type="ResNet50", bn_kwargs={}, num_classes=num_classes, include_top=False, ), neck=dict( type="ChannelMapperNeck", in_channels=[512, 1024, 2048], out_indices=[2, 3, 4], out_channel=256, kernel_size=1, extra_convs=1, norm_layer=dict( type=nn.BatchNorm2d, num_features=256, ), ), position_embedding=dict( type="PositionEmbeddingSine", num_pos_feats=128, temperature=10000, normalize=True, offset=-0.5, ), transformer=dict( type="DeformableDetrTransformer", encoder=dict( type="DeformableDetrTransformerEncoder", embed_dim=256, num_heads=8, feedforward_dim=1024, attn_dropout=0.1, ffn_dropout=0.1, num_layers=6, post_norm=False, num_feature_levels=4, ), decoder=dict( type="DeformableDetrTransformerDecoder", embed_dim=256, num_heads=8, feedforward_dim=1024, attn_dropout=0.1, ffn_dropout=0.1, num_layers=6, return_intermediate=True, num_feature_levels=4, as_two_stage=as_two_stage, with_box_refine=with_box_refine, ), num_feature_levels=4, two_stage_num_proposals=num_queries, as_two_stage=as_two_stage, ), embed_dim=256, num_classes=num_classes, num_queries=num_queries, aux_loss=True, post_process=dict( type="DeformDetrPostProcess", select_box_nums_for_evaluation=300 ), criterion=dict( type="DeformableCriterion", num_classes=num_classes, matcher=dict( type="HungarianMatcher", cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, use_focal=True, alpha=0.25, gamma=2.0, ), weight_dict={ "loss_class": 1.0, "loss_bbox": 5.0, "loss_giou": 2.0, }, loss_class_type="focal_loss", alpha=0.25, gamma=2.0, ), as_two_stage=as_two_stage, with_box_refine=with_box_refine, )

其中, model 下面的 type 表示定义的模型名称,剩余的变量表示模型的其他组成部分。这样定义模型的好处在于我们可以很方便的替换我们想要的结构。 例如,如果我们想训练一个backbone为efficientnet的模型,只需要将 model 下面的 backbone 替换掉,并相应设置 neck 就行。

数据增强

model 的定义一样,数据增强的流程是通过在config配置文件中定义 data_loaderval_data_loader 这两个dict来实现的,分别对应着训练集和验证集的处理流程。以 data_loader 为例:

data_loader = dict( type=torch.utils.data.DataLoader, dataset=dict( type="Coco", data_path="./tmp_data/mscoco/train_lmdb/", transforms=[ dict( type="RandomFlip", px=0.5, py=0, ), dict( type="RandomSelectOne", p=1.0, transforms=[ dict( type="Resize", img_scale=[ (480, 1333), (512, 1333), (544, 1333), (576, 1333), (608, 1333), (640, 1333), (672, 1333), (704, 1333), (736, 1333), (768, 1333), (800, 1333), ], multiscale_mode="value", keep_ratio=True, ), dict( type=Compose, transforms=[ dict( type="Resize", img_scale=[ (400, 9999999), (500, 9999999), (600, 9999999), ], multiscale_mode="value", keep_ratio=True, ), dict( type="RandomSizeCrop", min_size=384, max_size=600, filter_area=False, ), dict( type="Resize", img_scale=[ (480, 1333), (512, 1333), (544, 1333), (576, 1333), (608, 1333), (640, 1333), (672, 1333), (704, 1333), (736, 1333), (768, 1333), (800, 1333), ], multiscale_mode="value", keep_ratio=True, ), ], ), ], ), dict( type="ToTensor", to_yuv=True, use_yuv_v2=False, ), dict( type="Normalize", mean=128, std=128, ), ], ), sampler=dict(type=torch.utils.data.DistributedSampler), batch_size=batch_size_per_gpu, shuffle=True, num_workers=4, pin_memory=True, collate_fn=collate_2d_pad, )

其中type直接用的pytorch自带的接口torch.utils.data.DataLoader,表示的是将 batch_size 大小的图片组合到一起。 这里面唯一需要关注的可能是 dataset 这个变量, CocoFromLMDB 表示从lmdb数据集中读取图片,路径也就是我们在第一部分数据集准备中提到的路径。 transforms 下面包含着一系列的数据增强。 val_data_loader 则没有施加数据增强。 你也可以通过在 transforms 中插入新的dict实现自己希望的数据增强操作。

训练策略

为了训练一个精度高的模型,好的训练策略是必不可少的。对于每一个训练任务而言,相应的训练策略同样都定义在其中的config文件中,从 float_trainer 这个变量就可以看出来。

float_trainer = dict( type="distributed_data_parallel_trainer", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode="fuse_bn", converters=[ dict( type="LoadCheckpoint", checkpoint_path=( "./tmp_pretrained_models/resnet50_imagenet/float-checkpoint-best.pth.tar" ), allow_miss=True, ignore_extra=True, verbose=True, ), ], ), find_unused_parameters=False, data_loader=data_loader, optimizer=dict( type="custom_param_optimizer", optim_cls=torch.optim.AdamW, model=model, optim_cfgs=dict(lr=1e-4, weight_decay=1e-4), custom_param_mapper=dict( norm_types={"weight_decay": 0.0}, backbone={"lr": 1e-5} ), ), batch_processor=batch_processor, num_epochs=50, device=None, callbacks=[ stat_callback, loss_show_update, grad_callback, dict( type="StepDecayLrUpdater", lr_decay_id=[40], step_log_interval=100, ), bn_callback, val_callback, ckpt_callback, ], train_metrics=dict( type="LossShow", ), val_metrics=dict( type="COCODetectionMetric", ann_file="./tmp_data/mscoco/instances_val2017.json", ), )

float_trainer 从大局上定义了我们的训练方式,包括使用多卡分布式训练(distributed_data_parallel_trainer),模型训练的epoch次数,以及优化器的选择。 其中优化器是由 custom_param_optimizer 包裹的 AdamW 优化器,可以对模型中的backbone和norm layers进行更精细的优化参数控制。 同时 callbacks 中体现了模型在训练过程中使用到的小策略以及用户想实现的操作,包括学习率的变换方式(StepDecayLrUpdater),在训练过程中验证模型的指标(Validation),以及保存(Checkpoint)模型的操作。 当然,如果你有自己希望模型在训练过程中实现的操作,也可以按照这种dict的方式添加。

注解

如果需要复现精度,config中的训练策略最好不要修改。否则可能会有意外的训练情况出现。

通过上面的介绍,你应该对config文件的功能有了一个比较清楚的认识。然后通过前面提到的训练脚本,就可以训练一个高精度的纯浮点的检测模型。 当然训练一个好的检测模型不是我们最终的目的,它只是做为一个pretrain为我们后面训练定点模型服务的。

量化模型训练

当我们有了纯浮点模型之后,就可以开始训练相应的定点模型了。和浮点训练的方式一样,我们只需要通过运行下面的脚本就可以训练定点模型了。

python3 tools/train.py --stage calibration --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py python3 tools/train.py --stage qat --config configs/detection/deform_detr/deform_detr_resnet50_mscoco.py

可以看到,我们的配置文件没有改变,只改变了 stage 的类型。 Deformable DETR在量化训练过程中建议加上calibration的流程。calibration可以为QAT的量化训练提供一个更好的初始化参数。 此时我们使用的训练策略来自于config文件中的calibration_trainer和qat_trainer。

calibration_batch_processor = copy.deepcopy(val_batch_processor) calibration_step = 100 calibration_trainer = dict( type="Calibrator", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode="fuse_bn", qconfig_params=dict( activation_calibration_observer="mse", ), converters=[ dict( type="LoadCheckpoint", checkpoint_path=os.path.join( ckpt_dir, "float-checkpoint-best.pth.tar" ), ), dict(type="Float2Calibration", convert_mode=convert_mode), ], ), data_loader=calibration_data_loader, batch_processor=calibration_batch_processor, num_steps=calibration_step, device=None, callbacks=[ stat_callback, val_callback, ckpt_callback, ], val_metrics=dict( type="COCODetectionMetric", ann_file="./tmp_data/mscoco/instances_val2017.json", ), log_interval=calibration_step / 10, )

我们首先通过calibration为QAT获取一个好的初始化参数,calibration加载训练好的浮点模型,然后采用 mse 的校准方式搜索量化参数, 其中 Float2Calibration 将模型由浮点模型转为校准模型。

qat_trainer = dict( type="distributed_data_parallel_trainer", model=model, model_convert_pipeline=dict( type="ModelConvertPipeline", qat_mode="fuse_bn", qconfig_params=dict( activation_qat_qkwargs=dict( averaging_constant=0, ), weight_qat_qkwargs=dict( averaging_constant=1, ), ), converters=[ dict(type="Float2QAT", convert_mode=convert_mode), dict( type="LoadCheckpoint", checkpoint_path=os.path.join( ckpt_dir, "calibration-checkpoint-best.pth.tar" ), ), ], ), data_loader=qat_data_loader, optimizer=dict( type="custom_param_optimizer", optim_cls=torch.optim.AdamW, model=model, optim_cfgs=dict(lr=1e-5, weight_decay=0), custom_param_mapper=dict( norm_types={"weight_decay": 0.0}, ), ), batch_processor=batch_processor, num_epochs=6, device=None, callbacks=[ stat_callback, loss_show_update, grad_callback, dict( type="StepDecayLrUpdater", lr_decay_id=[5], step_log_interval=500, ), val_callback, ckpt_callback, ], train_metrics=dict( type="LossShow", ), val_metrics=dict( type="COCODetectionMetric", ann_file="./tmp_data/mscoco/instances_val2017.json", ), )

接着,我们在calibration模型的基础上进行QAT量化训练。初始学习率设置为浮点训练的十分之一,训练的epoch次数也大大减少。 注意到 converter 中的 Float2QAT 将模型由浮点模型转换为QAT模型,再加载的calibration模型权重。 经过量化训练,Deformable DETR量化模型的精度可以达到浮点模型精度的99%以上。