汽车关键点检测模型训练
这篇教程主要是告诉大家如何利用HAT在汽车关键点数据集 CarFusion 上从头开始训练一个关键点检测模型，包括浮点，量化和定点模型。
CarFusion 是一个汽车关键点数据集，数据集中标注了汽车的12个关键点，包含汽车前后左右的车轮轮心，车灯位置以及车顶角点。开始模型训练之前，第一步是准备好数据集，这里我们需要先在官网申请数据集，官网链接为：CarFusion，下拉可以找到数据集申请表单，提交表单后可以收到数据集下载链接邮件。下载解压缩后数据目录结构如下所示：
CarFusion
|-- train
|-- car_butler1
|-- car_butler2
|-- car_craig1
|-- car_craig2
|-- car_fifth1
|-- car_fifth2
|-- car_morewood1
|-- car_morewood2
|-- test
|-- car_penn1
|-- car_penn2
这里我们只考虑更简单的情形，即从检测好的汽车图片中检测关键点。因此，我们需要首先根据数据标注框从图片中裁剪出汽车。运行一下命令即可一键将数据转换为裁剪好的格式。
python3 tools/dataset_converters/gen_carfusion_data.py --src-data-path ${data-dir} --out-dir ${cropped-data-dir} --num-workers 10
得到的文件夹组织样式是这样的：
cropped-data-dir
|-- train
|-- car_butler1
|-- car_butler2
|-- car_craig1
|-- car_craig2
|-- car_fifth1
|-- car_fifth2
|-- car_morewood1
|-- car_morewood2
|-- test
|-- car_penn1
|-- car_penn2
|-- simple_anno
|-- keypoints_train.json
|-- keypoints_test.json
其中 keypoints_test.json 和 keypoints_train.json 分别包含测试集和训练集的标注。 每个样本以jsonge字典格式存储，样式为： "img_path": keypoints_list，其中keypoints_list形状为(12,3)，每行前两个元素为关键点坐标(x,y)，第三个元素为关键点是否有效的标注，当关键点不在图内或无效时，第三个元素为0。
接着，可以将数据打包为 LMDB 格式的数据集。只需要运行下面的脚本，就可以成功实现转换：
python3 tools/datasets/carfusion_packer.py --src-data-dir ${cropped-data-dir} --target-data-dir ${pack-data-dir} --split-name train --pack-type lmdb --num-workers 10
python3 tools/datasets/carfusion_packer.py --src-data-dir ${cropped-data-dir} --target-data-dir ${pack-data-dir} --split-name test --pack-type lmdb --num-workers 10
打包完成之后，目录下的文件结构应该如下所示：
tmp_data
|-- carfusion
|-- test_lmdb
|-- train_lmdb
train_lmdb 和 val_lmdb 就是打包之后的训练数据集和验证数据集，接下来就可以开始训练模型。
模型训练
configs/keypoint/keypoint_efficientnetb0_carfusion.py 包含了本教程中模型训练相关的所有设置。
在网络开始训练之前，你可以使用以下命令先计算一下网络的计算量和参数数量：
python3 tools/calops.py --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
下一步就可以开始训练。训练可以通过下面的脚本来完成，在训练之前需要确认配置中数据集路径是否已经切换到已经打包好的数据集路径。
python3 tools/train.py --stage "float" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
python3 tools/train.py --stage "calibration" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
由于HAT算法包使用了注册机制，每一个训练任务都可以按照这种 train.py 加上 config 配置文件的形式启动。 train.py 是统一的训练脚本，与任务无关，我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的 config 配置文件里面。 上面的命令中 --stage 后面的参数可以是 "float" 、 "calibration"，分别可以完成浮点模型和量化模型的校准，其中量化模型的训练依赖于上一步浮点训练产出的浮点模型。
导出定点模型
完成量化训练后，便可以开始导出定点模型。可以通过下面命令来导出：
python3 tools/export_hbir.py --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
模型验证
在完成训练之后，可以得到训练完成的浮点、量化、定点模型。和训练方法类似，我们可以用相同方法来对训好的模型做指标验证，得到为 Float 、 QAT 的指标，分别为浮点、量化的指标。
python3 tools/predict.py --stage "float" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
python3 tools/predict.py --stage "calibration" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
和训练模型时类似， --stage 后面的参数为 "float" 和 "calibration"时， 分别可以完成对训练好的浮点模型、量化模型的验证。
对于定点模型精度验证，也可使用下面命令，但需要注意是必须要先导出hbir：
python3 tools/predict.py --stage "int_infer" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
模型推理
HAT提供了 infer_hbir.py 脚本对各阶段训练好的模型的推理结果进行可视化展示：
python3 tools/infer_hbir.py --config configs/keypoint/keypoint_efficientnetb0_carfusion.py --model-inputs img:${img1-path} --save-path ${save_path}
仿真上板精度验证
除了上述模型验证之外，我们还提供和上板完全一致的精度验证方法，可以通过下面的方式完成：
python3 tools/validation_hbir.py --stage "align_bpu" --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
定点模型检查和编译
在HAT中集成的量化训练工具链主要是为了地平线的计算平台准备的，因此，对于量化模型的检查和编译是必须的。我们在HAT中提供了模型检查的接口，可以让用户定义好量化模型之后，先检查能否在 BPU 上正常运行：
python3 tools/model_checker.py --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
在模型训练完成后，可以通过 compile_perf_hbir 脚本将量化模型编译成可以上板运行的 hbm 文件，同时该工具也能预估在 BPU 上的运行性能：
python3 tools/compile_perf_hbir.py --config configs/keypoint/keypoint_efficientnetb0_carfusion.py
以上就是从数据准备到生成量化可部署模型的全过程。
训练细节
在这个说明中，我们对模型训练需要注意的一些事项进行说明，主要为 config 的一些相关设置。
模型构建
针对轻量的汽车关键点检测任务，网络模型 HeatmapKeypointModel 使用efficienet-b0作为backbone，加上三层转置卷积层生成热力图，然后从热力图解码关键点坐标。我们通过在 config 配置文件中定义 model 这样的一个 dict 型变量，就可以方便的实现对模型的定义和修改。
model = dict(
type="HeatmapKeypointModel",
backbone=dict(
type="efficientnet",
model_type="b0",
num_classes=1,
bn_kwargs={},
activation="relu",
use_se_block=False,
include_top=False,
),
decode_head=dict(
type="DeconvDecoder",
in_channels=320,
out_channels=NUM_LDMK,
input_index=4,
num_conv_layers=3,
num_deconv_filters=[128, 128, 128],
num_deconv_kernels=[4, 4, 4],
final_conv_kernel=3,
),
loss=dict(type="MSELoss", reduction="mean"),
post_process=dict(
type="HeatmapDecoder",
scale=4,
mode="averaged",
),
)
模型包含 backbone ，转置卷积构成的 decode_head ， losses 和 post_process 模块，在HeatmapKeypointModel中, 输入为裁剪好的汽车图片，backbone 用于提取图像特征， decoder 上采样输出 heatmap ， losses 使用根据heatmap位置加权的 MSELoss 来作为训练的 loss ，post_process 使用 HeatmapDecoder 将heatmap输出转换为关键点位置预测值。
数据增强
跟 model 的定义一样，数据增强的流程是通过在 config 配置文件中定义 data_loader 和 val_data_loader 这两个 dict 来实现的，分别对应着训练集和验证集的处理流程。以 data_loader 为例，数据增强使用了 RandomFlip 、 Resize 、 RandomPadLdmkData 和 GaussianNoise。对于关键点检测任务，还需要使用 GenerateHeatmapTarget 从关键点标注生成热力图(Heatmap)目标。
data_loader = dict(
type=torch.utils.data.DataLoader,
dataset=dict(
type="CarfusionPackData",
data_path=f"{data_root}/train_lmdb",
transforms=[
dict(type="RandomFlip", px=0.5),
dict(
type="Resize",
img_scale=image_size,
keep_ratio=True,
),
dict(
type="RandomPadLdmkData",
size=image_size,
),
dict(
type="AddGaussianNoise",
prob=0.2,
mean=0,
sigma=2,
),
dict(
type="GenerateHeatmapTarget",
num_ldmk=NUM_LDMK,
feat_stride=4,
heatmap_shape=(32, 32),
sigma=1.0,
),
dict(
type="ToTensor",
to_yuv=True,
use_yuv_v2=False,
),
dict(
type="Normalize",
mean=128.0,
std=128.0,
),
],
),
sampler=dict(type=torch.utils.data.DistributedSampler),
batch_size=batch_size_per_gpu,
shuffle=True,
num_workers=8,
pin_memory=True,
collate_fn=hat.data.collates.collate_2d,
)
因为最终跑在 BPU 上的模型使用的是 YUV444 的图像输入，而一般的训练图像输入都采用 RGB 的形式， 所以HAT在 ToTensor 数据transform中提供了 to_yuv=True 的选项，将 RGB 格式图片转换为 YUV444 的格式。 HAT还提供了 batch_processor 接口对数据进行批处理，但此处没有添加额外的增强处理。其中 loss_collector 是一个获取当前批量数据的 loss 的函数，由于模型返回值为 (pred, loss) 的 Tuple，因此取模型输出tuple的index=1的值，即为 loss 。
batch_processor = dict(
type="MultiBatchProcessor",
need_grad_update=True,
loss_collector=collect_loss_by_index(1),
)
验证集的数据转换相对简单很多，如下所示：
val_data_loader = dict(
type=torch.utils.data.DataLoader,
dataset=dict(
type="CarfusionPackData",
data_path=f"{data_root}/test_lmdb",
transforms=[
dict(
type="Resize",
img_scale=image_size,
keep_ratio=True,
),
dict(
type="RandomPadLdmkData",
size=image_size,
random=False,
),
dict(
type="ToTensor",
to_yuv=True,
use_yuv_v2=False,
),
dict(
type="Normalize",
mean=128.0,
std=128.0,
),
],
),
batch_size=batch_size_per_gpu,
sampler=dict(type=torch.utils.data.DistributedSampler),
shuffle=False,
num_workers=8,
pin_memory=True,
collate_fn=hat.data.collates.collate_2d,
)
val_batch_processor = dict(
type="MultiBatchProcessor",
need_grad_update=False,
)
训练策略
configs/keypoint/keypoint_efficientnetb0_carfusion.py 文件中的 float_trainer ， calibration_trainer 分别对应浮点与量化模型的训练策略。下面以 float_trainer 训练策略示例：
float_trainer = dict(
type="distributed_data_parallel_trainer",
model=model,
model_convert_pipeline=dict(
type="ModelConvertPipeline",
converters=[
dict(
type="LoadCheckpoint",
checkpoint_path=pretrain_model_path,
allow_miss=True,
ignore_extra=True,
),
],
),
data_loader=data_loader,
optimizer=dict(type=torch.optim.AdamW, lr=0.001, weight_decay=5e-2),
batch_processor=batch_processor,
num_epochs=30,
device=None,
callbacks=[
stat_callback,
loss_show_update,
dict(
type="CosLrUpdater",
warmup_len=0,
warmup_by="epoch",
step_log_interval=100,
),
val_callback,
ckpt_callback,
],
train_metrics=[
dict(
type="LossShow",
),
],
sync_bn=True,
val_metrics=[
dict(type="PCKMetric", alpha=0.1, feat_stride=4, img_shape=image_size),
dict(
type="MeanKeypointDist",
feat_stride=4,
),
],
)
float_trainer 从大局上定义了我们的训练方式，包括使用多卡分布式训练（distributed_data_parallel_trainer），模型训练的epoch次数，以及优化器的选择。其中 model_convert_pipeline 定义了模型开始训练前的转换操作，例如此处模型首先载入 Efficientnet-b0 在 ImageNet 上的预训练模型。 optimizer 此处使用的是 AdamW 优化器，配合 lr=0.001 的学习率和 weight_decay=0.05 。实验表明较大的weight_decay值能够帮助训练量化性能更好的浮点模型。
callbacks 定义了模型在训练过程中使用到的小策略以及用户想实现的操作，包括学习率的变换方式 CosLrUpdater，在训练过程中验证模型的指标(Validation)，以及保存(Checkpoint)模型的操作。当然，如果你有自己希望模型在训练过程中实现的操作，也可以按照这种dict的方式添加。
train_metrics 和 val_metrics 分别定义了模型训练和验证时监测的指标(Metric)。train_metrics 中的 LossShow 可以监测训练损失。 val_metrics 中的 PCKMetric 计算了在一定距离阈值内与相应真实关键点重合的关键点的百分比；MeanKeypointDist 计算了关键点预测点与真实关键点的平均距离误差。
总之 float_trainer 负责将模型整个浮点训练的逻辑给串联起来。
量化模型校准
当我们有了纯浮点模型后，我们首先通过校准（Calibration）操作使用一些数据样本计算模型量化各层的缩放（scale）参数，从而对模型进行int8量化。相关的配置为：
calibration_data_loader = copy.deepcopy(data_loader)
calibration_data_loader.pop("sampler") # Calibration do not support DDP or DP
calibration_data_loader["batch_size"] = batch_size_per_gpu * 4
calibration_data_loader["dataset"]["transforms"] = val_data_loader["dataset"][
"transforms"
]
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="percentile",
activation_calibration_qkwargs=dict(
percentile=99.975,
),
),
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="PCKMetric", alpha=0.1, feat_stride=4, img_shape=image_size),
dict(
type="MeanKeypointDist",
feat_stride=4,
),
],
log_interval=calibration_step / 10,
)
模型检查编译和仿真上板精度验证
对于HAT来说，量化模型的意义在于可以在BPU上直接运行。因此，对于量化模型的检查和编译是必须的。
前文提到的 compile_perf_hbir 脚本也可以让用户定义好量化模型之后，先检查能否在BPU上正常运行， 并可通过 validation_hbir 脚本获取模型上板精度。用法同前文。
