10.4.15. Motr多目标跟踪模型训练

这篇教程主要是告诉大家如何利用HAT在目标跟踪数据集 MOT17 上从头开始训练一个 Motr 模型，包括浮点、量化和定点模型。

MOT17 是多目标跟踪中用的比较多的数据集，很多先进的多目标跟踪研究都会优先基于这个数据集做好验证。 开始训练模型之前，第一步是准备好数据集，这里我们下载官方的数据集以及相应的标签数据 MOT17DATASET 。

解压缩之后数据目录结构如下所示：

tmp_data
  |-- train
      |-- MOT17-02-DPM
          |-- det
              |-- det.txt
          |-- gt
              |-- gt.txt
          |-- img1
              |-- 000001.jpg
              |-- 000002.jpg
              |-- ...
          |-- seqinfo.ini
      |-- MOT17-02-FRCNN
      |-- MOT17-02-SDP
      |-- MOT17-04-DPM
      |-- MOT17-04-FRCNN
      |-- MOT17-04-SDP
      |-- ...
  |-- test
      |-- MOT17-01-DPM
      |-- ...

其中 MOT17-02-DPM 视频序列的名字， det.txt 为检测器的结果， 比如 MOT17-02-DPM 文件夹中的就是使用 DPM 算法检测的结果， gt 下的 gt.txt 为标签数据， img1 中为图像数据。

10.4.15.1. 训练流程

如果你只是想简单的把 Motr 的模型训练起来，那么可以首先阅读一下这一章的内容。 和其他任务一样，对于所有的训练，评测任务，HAT统一采用 tools + config 的形式来完成。 在准备好原始数据集之后，可以通过下面的流程，方便地完成整个训练的流程。

10.4.15.1.1. 数据集准备

由于官方的测试集无gt，因此我们将训练集拆成一半，每个视频前一半帧作为训练集，后一半帧作为验证集。 HAT 提供脚本将训练集拆分，只需要运行下面的脚本：

python3 tools/dataset_converters/gen_mot_data.py --src-data-path ${data-dir} --out-dir ${data-dir}/split_data

运行完上面脚本后，将会得到下面生成类似下面结构的文件夹：

tmp_data
  |-- train
  |-- test
  |-- split_data
      |-- train
          |-- MOT17-02-SDP
              |-- gt
                  |-- gt.txt
              |-- labels_with_ids
                  |-- 000001.txt
                  |-- 000002.txt
              |-- img1
                  |-- 000001.jpg
                  |-- 000002.jpg
                  |-- ...
          |-- MOT17-04-SDP
          |-- MOT17-05-SDP
          |-- MOT17-09-SDP
          |-- MOT17-10-SDP
          |-- MOT17-11-SDP
          |-- MOT17-13-SDP
      |-- test
          |-- MOT17-02-SDP
              |-- gt
                  |-- gt.txt
              |-- labels_with_ids
                  |-- 000301.txt
                  |-- 000302.txt
              |-- img1
                  |-- 000301.jpg
                  |-- 000302.jpg
                  |-- ...
          |-- MOT17-04-SDP
          |-- MOT17-05-SDP
          |-- MOT17-09-SDP
          |-- MOT17-10-SDP
          |-- MOT17-11-SDP
          |-- MOT17-13-SDP

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

python3 tools/datasets/mot_packer.py --src-data-dir ${data-dir}/split_data --split-name train --pack-type lmdb  --num-workers 10 --target-data-dir ${target-data-dir}
python3 tools/datasets/mot_packer.py --src-data-dir ${data-dir}/split_data --split-name test --pack-type lmdb  --num-workers 10 --target-data-dir ${target-data-dir}

上面这两条命令分别对应转换训练数据集和验证数据集。

评测精度时，需要用到验证数据集的标签数据，因此我们做一个软连接，如下所示：

ln -s  ${data-dir}/split_data/test/ ${target-data-dir}/test_gt

打包和软连接完成之后， ${target-data-dir} 目录下的文件结构应该如下所示：

${target-data-dir}
  |-- train_lmdb
  |-- test_lmdb
  |-- test_gt

train_lmdb 和 test_lmdb 就是打包之后的训练数据集和验证数据集， test_gt 里面是验证集的标签数据，接下来就可以开始训练模型。

10.4.15.1.2. 模型训练

由于 Motr 模型中的 qim 模块输入依赖于一些后处理，因此我们将整个 Motr 模型拆分到两个config中。

除去生成定点模型、编译、模型checker、模型计算量需要使用 qim 的 config，其余均使用 motr 的 config 即可，详细使用情况可见下面章节。

以下描述中，第一个模型为 Motr 的基础模块，第二个模型为 qim 模块，未特别注明则是两个模块串联的模型。

在网络开始训练之前，你可以使用以下命令先计算一下网络的计算量和参数数量：

python3 tools/calops.py --config configs/track_pred/motr_efficientnetb3_mot17.py  # 第一个模型的计算量
python3 tools/calops.py --config configs/track_pred/motr_efficientnetb3_mot17_qim.py  # 第二个模型的计算量

下一步就可以开始训练。训练也可以通过下面的脚本来完成，在训练之前需要确认配置中数据集路径是否已经切换到已经打包好的数据集路径。

python3 tools/train.py --stage "float" --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/train.py --stage "calibration" --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/train.py --stage "qat" --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/train.py --stage "int_infer" --config configs/track_pred/motr_efficientnetb3_mot17.py # 第一个模型
python3 tools/train.py --stage "int_infer" --config configs/track_pred/motr_efficientnetb3_mot17_qim.py # 第二个模型

由于HAT算法包使用了注册机制，使得每一个训练任务都可以按照这种 train.py 加上 config 配置文件的形式启动。 train.py 是统一的训练脚本，与任务无关，我们需要训练什么样的任务、使用什么样的数据集以及训练相关的超参数设置都在指定的 config 配置文件里面。 上面的命令中 --stage 后面的参数可以是 "float" 、 "calibration" 、 "qat" 、 "int_infer"， 分别可以完成浮点模型、量化模型的训练以及量化模型到定点模型的转化， 其中量化模型的训练依赖于上一步浮点训练产出的浮点模型，定点模型的转化依赖于量化训练产生的量化模型。

对于该模型，量化模型到定点模型的转化需要使用 qim 模块的 config。

10.4.15.1.3. 模型验证

在完成训练之后，可以得到训练完成的浮点、量化或定点模型。和训练方法类似， 我们可以用相同方法来对训好的模型做指标验证，得到为 Float 、 Calibration 、 QAT 和 Quantized 的指标，分别为浮点、量化和完全定点的指标。

python3 tools/predict.py --stage "float" --config configs/track_pred/motr_efficientnetb3_mot17.py

python3 tools/predict.py --stage "calibration" --config configs/track_pred/motr_efficientnetb3_mot17.py

python3 tools/predict.py --stage "qat" --config configs/track_pred/motr_efficientnetb3_mot17.py

python3 tools/predict.py --stage "int_infer" --config configs/track_pred/motr_efficientnetb3_mot17.py

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

10.4.15.1.4. 模型推理

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

python3 tools/infer.py --config configs/track_pred/motr_efficientnetb3_mot17.py --model-inputs imagedir:${img-dir},num:${num_pic} --save-path ${save_path}

10.4.15.1.5. 仿真上板精度验证

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

python3 tools/align_bpu_validation.py --config configs/track_pred/motr_efficientnetb3_mot17.py

10.4.15.1.6. 定点模型检查和编译

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

python3 tools/model_checker.py --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/model_checker.py --config configs/track_pred/motr_efficientnetb3_mot17_qim.py

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

python3 tools/compile_perf.py --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/compile_perf.py --config configs/track_pred/motr_efficientnetb3_mot17_qim.py

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

10.4.15.1.7. ONNX模型导出

如果想要导出onnx模型, 运行下面的命令即可:

python3 tools/export_onnx.py --config configs/track_pred/motr_efficientnetb3_mot17.py
python3 tools/export_onnx.py --config configs/track_pred/motr_efficientnetb3_mot17_qim.py

10.4.15.2. 训练细节

在这个说明中，我们对模型训练需要注意的一些事项进行说明，主要为 config 的一些相关设置。

10.4.15.2.1. 模型构建

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

num_queries = 256
num_classes = 1
model = dict(
    type="Motr",
    backbone=dict(
        type="efficientnet",
        bn_kwargs={},
        model_type="b3",
        num_classes=1000,
        include_top=False,
        activation="relu",
        use_se_block=False,
    ),
    head=dict(
        type="MotrHead",
        transformer=dict(
            type="MotrDeformableTransformer",
            pos_embed=dict(
                type="PositionEmbeddingSine",
                num_pos_feats=128,
                normalize=True,
                temperature=20,
            ),
            d_model=256,
            num_queries=num_queries,
            dim_feedforward=1024,
            dropout=0.0,
            return_intermediate_dec=True,
            extra_track_attn=True,
            enc_n_points=1,
            dec_n_points=1,
        ),
        num_classes=num_classes,
        in_channels=[384],
        max_per_img=num_queries,
    ),
    criterion=dict(
        type="MotrCriterion",
        num_classes=num_classes,
    ),
    post_process=dict(
        type="MotrPostProcess",
    ),
    track_embed=dict(
        type="QueryInteractionModule",
        dim_in=256,
        hidden_dim=1024,
    ),
)

模型除了 backbone 之外，还有 head、 criterion、 post_process 和 track_embed 模块， 在 Motr 中, backbone 主要是提取图像的特征， head 主要是由特征来得到预测的类别、位置和特征。 criterion 是训练时计算 loss 的模块, post_process 主要是后处理部分， track_embed 是用来更新已跟踪上的目标query的模块(即 qim 模块)。

10.4.15.2.2. 数据增强

跟 model 的定义一样，数据增强的流程是通过在 config 配置文件中定义 data_loader 和 val_data_loader 这两个 dict 来实现的， 分别对应着训练集和验证集的处理流程。以 data_loader 为例， 数据增强使用了 SeqRandomFlip、 RandomSelectOne、 SeqResize、 SeqRandomSizeCrop、 SeqToTensor、 SeqNormalize 来增加训练数据的多样性，增强模型的泛化能力。

因为最终跑在 BPU 上的模型使用的是 YUV444 的图像输入，而一般的训练图像输入都采用 RGB 的形式，所以HAT提供 SeqBgrToYuv444 的数据增强来将 RGB 转到 YUV444 的格式。

from torchvision.transforms import Compose
train_lmdb="./tmp_data/mot17/train_lmdb"
train_batch_size_per_gpu = 1
data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="Mot17Dataset",
        data_path=train_lmdb,
        sampler_lengths=[2,3,4,5],
        sample_mode="random_interval",
        sample_interval=10,
        sampler_steps=[100, 180, 240],
        transforms=[
            dict(
                type="SeqRandomFlip",
                px=0.5,
                py=0,
            ),
            dict(
                type="RandomSelectOne",
                transforms=[
                    dict(
                        type="SeqResize",
                        img_scale=[
                            (608, 1536),
                            (640, 1536),
                            (672, 1536),
                            (704, 1536),
                            (736, 1536),
                            (768, 1536),
                            (800, 1536),
                            (832, 1536),
                            (864, 1536),
                            (896, 1536),
                            (928, 1536),
                            (960, 1536),
                            (992, 1536),
                        ],
                        multiscale_mode="value",
                        keep_ratio=True,
                        rm_neg_coords=False,
                        divisor=2,
                    ),
                    dict(
                        type=Compose,
                        transforms=[
                            dict(
                                type="SeqResize",
                                img_scale=[
                                    (400, 9999999),
                                    (500, 9999999),
                                    (600, 9999999),
                                ],
                                multiscale_mode="value",
                                keep_ratio=True,
                                rm_neg_coords=False,
                                divisor=2,
                            ),
                            dict(
                                type="SeqRandomSizeCrop",
                                min_size=384,
                                max_size=600,
                                filter_area=False,
                                rm_neg_coords=False,
                            ),
                            dict(
                                type="SeqResize",
                                img_scale=[
                                    (608, 1536),
                                    (640, 1536),
                                    (672, 1536),
                                    (704, 1536),
                                    (736, 1536),
                                    (768, 1536),
                                    (800, 1536),
                                    (832, 1536),
                                    (864, 1536),
                                    (896, 1536),
                                    (928, 1536),
                                    (960, 1536),
                                    (992, 1536),
                                ],
                                multiscale_mode="value",
                                keep_ratio=True,
                                rm_neg_coords=False,
                                divisor=2,
                            ),
                        ],
                    ),
                ],
                p=1,
            ),
            dict(
                type="SeqToTensor",
                to_yuv=False,
            ),
            dict(type="SeqBgrToYuv444", rgb_input=True),
            dict(
                type="SeqNormalize",
                mean=128.0,
                std=128.0,
            ),
        ]
    ),
    sampler=dict(type="DistSetEpochDatasetSampler"),
    batch_size=train_batch_size_per_gpu,
    pin_memory=True,
    shuffle=True,
    num_workers=2,
    collate_fn=hat.data.collates.collate_mot_seq,
)

batch_processor 中传入一个 loss_collector 函数，用于获取当前批量数据的 loss，如下所示：

def loss_collector(outputs: dict):
    losses = []
    for loss_name, loss in outputs.items():
        mean_loss = sum(loss)/len(loss)
        losses.append(mean_loss)
    return losses

batch_processor = dict(
    type="MultiBatchProcessor",
    need_grad_update=True,
    loss_collector=loss_collector,
)

验证集的数据转换相对简单很多，如下所示：

import torch
val_lmdb = "./tmp_data/mot17/test_lmdb"
test_batch_size_per_gpu = 1
val_data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="Mot17Dataset",
        data_path=val_lmdb,
        sampler_lengths=[1],
        sample_mode="fixed_interval",
        sample_interval=10,
        transforms=[
            dict(
                type="SeqResize",
                img_scale=(800, 1422),
                keep_ratio=False,
            ),
            dict(
                type="SeqToTensor",
                to_yuv=False,
            ),
            dict(type="SeqBgrToYuv444", rgb_input=True),
            dict(
                type="SeqNormalize",
                mean=128.0,
                std=128.0,
            ),
        ]
    ),
    batch_size=test_batch_size_per_gpu,
    pin_memory=True,
    shuffle=False,
    num_workers=2,
    collate_fn=hat.data.collates.collate_mot_seq,
)

val_batch_processor = dict(
    type="MultiBatchProcessor",
    need_grad_update=False,
)

10.4.15.2.3. 训练策略

在 Mot17 数据集上训练浮点模型使用 Stepdecay 的学习策略, 以及对 weight 的参数施加 L2 norm。 configs/track_pred/motr_efficientnetb3_mot17.py 文件中的 float_trainer, calibration_trainer, qat_trainer, int_trainer 分别对应浮点、量化、定点模型的训练策略。 下面为 float_trainer 训练策略示例：

import torch

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/fcos_efficientnetb3_mscoco/float-checkpoint-best.pth.tar"
                ),
                allow_miss=True,
                ignore_extra=True,
                verbose=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        params={
            "backbone": dict(lr=2e-5),
            "sampling_offsets": dict(lr=2e-5),
        },
        eps=1e-8,
        betas=(0.9, 0.999),
        lr=2e-4,
        weight_decay=1e-4,
    ),
    batch_processor=batch_processor,
    num_epochs=400,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        grad_callback,
        dict(
            type="StepDecayLrUpdater",
            lr_decay_id=[200],
            step_log_interval=500,
        ),
        bn_callback,
        val_callback,
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="MotMetric",
        gt_dir=val_gt,
        save_prefix=ckpt_dir+"/metric_out/float_train",
    ),
    sync_bn=True,
)

10.4.15.2.4. 量化训练

关于量化训练中的关键步骤，比如准备浮点模型、算子替换、插入量化和反量化节点、设置量化参数以及算子的融合等，请阅读 量化感知训练 章节的内容。 这里主要讲一下 HAT 的多目标跟踪模型中如何定义和使用量化模型。

在模型准备的好情况下，包括量化已有的一些模块完成之后，HAT在训练脚本中统一使用下面的脚本将浮点模型映射到定点模型上来。

model.fuse_model()
model.set_qconfig()
horizon.quantization.prepare_qat(model, inplace=True)

量化训练的整体策略可以直接沿用浮点训练的策略，但学习率和训练长度需要适当调整。 因为有浮点预训练模型，所以量化训练的学习率 Lr 可以很小，一般可以从 0.001 或 0.0001 开始，并可以搭配 StepLrUpdater 做1-2次 scale=0.1 的 Lr 调整； 同时训练的长度不用很长。此外 weight decay 也会对训练结果有一定影响。

Motr 示例模型的量化训练策略可见 configs/track_pred/motr_efficientnetb3_mot17.py 文件。

10.4.15.2.5. 模型检查编译和仿真上板精度验证

对于HAT来说，量化模型的意义在于可以在 BPU 上直接运行。因此，对于量化模型的检查和编译是必须的。 前文提到的 compile_perf 脚本也可以让用户定义好量化模型之后，先检查能否在 BPU 上正常运行，并可通过 align_bpu_validation 脚本获取模型上板精度。用法同前文。