QCNet Trajectory Prediction Model Training

This tutorial primarily guides you on training a QCNet model from scratch on the Argoverse 2 dataset, including floating-point, quantized, and fixed-point models. QCNet is a trajectory prediction model, which can be referenced in the paperQuery-Centric Trajectory Prediction.

Training Process

If you want to quickly train the QCNet model, start by reading this section. Similar to other tasks, all training and evaluation tasks in HAT are accomplished using a tools + config approach. After preparing the raw dataset, you can conveniently complete the entire training process by following the steps below.

Dataset Preparation

The first step before training the model is to prepare the dataset. You can download it from the Argoverse 2 Dataset. Download the following: Training, Validation, and Test.

After downloading, extract and organize the folders as follows:

data
  |-- argoverse-2
      |-- train
      |-- val
      |-- test

To speed up the training process, we packed the dataset information files into lmdb format. You can convert the dataset by running the following script:

python3 tools/datasets/argoverse2_packer.py --src-data-dir source-data-dir --mode train --pack-type lmdb  --num-workers 10 --target-data-dir target-data-dir
python3 tools/datasets/argoverse2_packer.py --src-data-dir source-data-dir --mode val --pack-type lmdb  --num-workers 10 --target-data-dir target-data-dir

These commands respectively convert the training and validation datasets. After packing, the target-data-dir directory structure should be as follows:

tmp_data
    |-- argoverse-2
        |-- qc_data
            |-- train_lmdb
            |-- val_lmdb

train_lmdb and val_lmdb are the packed training and validation datasets. Now you can start training the model.

Model Training

Next, you can start training. Training can be done using the following scripts. Before training, ensure that the dataset path in the configuration has been switched to the packed dataset path.

python3 tools/train.py --stage "float" --config configs/traj_pred/qcnet_oe_argoverse2.py
python3 tools/train.py --stage "calibration" --config configs/traj_pred/qcnet_oe_argoverse2.py
python3 tools/train.py --stage "qat" --config configs/traj_pred/qcnet_oe_argoverse2.py

These commands respectively complete the training of the floating-point model and the fixed-point model. Training the fixed-point model requires a pre-trained floating-point model. For more details, please read the section on Quantized Awareness Training.

Export the Fixed-Point Model

After completing quantization training, you can start exporting the fixed-point model. You can export it using the following command:

python3 tools/export_hbir.py --config configs/traj_pred/qcnet_oe_argoverse2.py

Due to the relative spatiotemporal encoding used in QCNet, which introduces spatiotemporal invariance, we can reuse the features of historical frame encodings during model deployment. We offer a cold start model deployment solution: encoding frames sequentially, where the cold start is based on QAT training of the floating-point model from the hot start.

Model Validation

After completing the training, we can obtain the trained floating-point and quantized models. Similar to the training method, we can use the same approach to evaluate the metrics of the trained models. This will yield metrics for Float, Calibration, and Qat. The former provides the metrics for the floating-point model, while the latter two provide the metrics for the models obtained through quantization calibration and quantized awareness training, respectively.

python3 tools/predict.py --stage "float" --config configs/traj_pred/qcnet_oe_argoverse2.py
python3 tools/predict.py --stage "calibration" --config configs/traj_pred/qcnet_oe_argoverse2.py
python3 tools/predict.py --stage "qat" --config configs/traj_pred/qcnet_oe_argoverse2.py

For fixed-point model accuracy validation, use the following command, but note that it must be exported as hbir first:

python3 tools/predict.py --stage "int_infer" --config configs/traj_pred/qcnet_oe_argoverse2.py

Model Inference and Results Visualization

HAT provides the infer_hbir.py script to visualize the inference results of the fixed-point model:

python3 tools/infer_hbir.py --config configs/traj_pred/qcnet_oe_argoverse2.py --model-inputs ${model_input} --save-path ${save_path} --use-dataset

python3 tools/infer_hbir.py --config configs/traj_pred/qcnet_oe_argoverse2.py --model-inputs ${model_input} --save-path ${save_path}

Fixed-Point Model Checking and Compilation

The quantization training toolchain integrated in HAT is mainly prepared for the Horizon computing platform, so it is necessary to check and compile the quantized models. We provide a model checking interface in HAT, which can be used to check whether the defined quantization model can run normally on the BPU before training:

python3 tools/model_checker.py --config configs/traj_pred/qcnet_oe_argoverse2.py

After training the model, you can compile the quantized model into an hbm file that can be run on board using the compile_perf_hbir script. This tool can also estimate the performance on BPU:

python3 tools/compile_perf_hbir.py --config configs/traj_pred/qcnet_oe_argoverse2.py

The compiled hbm file is stored at ${ckpt_dir}/compile.

This completes the entire process from data preparation to generating a deployable quantized model.

Training Details

In this note, we explain some things that need to be considered for model training, mainly including settings related to config.

Model Construction

The QCNet network structure can be referenced from the paper. We define and modify the model conveniently by defining a dict variable named model in the config file.


model = dict(
    type="QCNetOE",
    encoder=dict(
        type="QCNetOEEncoder",
        map_encoder=dict(
            type="QCNetOEMapEncoder",
            input_dim=input_dim,
            hidden_dim=hidden_dim,
            num_historical_steps=num_historical_steps,
            num_freq_bands=num_freq_bands,
            num_layers=num_map_layers,
            num_heads=num_heads,
            head_dim=head_dim,
            dropout=dropout,
        ),
        agent_encoder=dict(
            type="QCNetOEAgentEncoderStream",
            input_dim=input_dim,
            hidden_dim=hidden_dim,
            num_historical_steps=num_historical_steps,
            time_span=time_span,
            num_freq_bands=num_freq_bands,
            num_layers=num_agent_layers,
            num_heads=num_heads,
            head_dim=head_dim,
            num_t2m_steps=num_t2m_steps,
            num_pl2a=num_pl2a,
            num_a2a=num_a2a,
            dropout=dropout,
            save_memory=True,
            stream_infer=True,
            reuse_agent_rembs=reuse_agent_rembs,
            deploy=False,
            quant_infer_cold_start=quant_infer_cold_start,
        ),
    ),
    decoder=dict(
        type="QCNetOEDecoder",
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        output_dim=output_dim,
        num_historical_steps=num_historical_steps,
        num_future_steps=num_future_steps,
        num_modes=num_modes,
        num_recurrent_steps=num_recurrent_steps,
        num_t2m_steps=num_t2m_steps,
        num_freq_bands=num_freq_bands,
        num_layers=num_dec_layers,
        num_heads=num_heads,
        head_dim=head_dim,
        dropout=dropout,
        split_rec_modules=split_rec_modules,
        reuse_agent_rembs=reuse_agent_rembs,
        deploy=False,
        quant_infer_cold_start=quant_infer_cold_start,
    ),
    loss=dict(
        type="QCNetOELoss",
        output_dim=output_dim,
        num_historical_steps=num_historical_steps,
        num_future_steps=num_future_steps,
    ),
    preprocess=dict(
        type="QCNetOEPreprocess",
        input_dim=input_dim,
        num_historical_steps=num_historical_steps,
        time_span=time_span,
        num_t2m_steps=num_t2m_steps,
        num_agent_layers=num_agent_layers,
        num_pl2a=num_pl2a,
        num_a2a=num_a2a,
        stream=True,
        save_memory=True,
        deploy=False,
        quant_infer_cold_start=quant_infer_cold_start,
    ),
    postprocess=dict(
        type="QCNetOEPostprocess",
        output_dim=output_dim,
        num_historical_steps=num_historical_steps,
    ),
    quant_infer_cold_start=quant_infer_cold_start,
)

Data Loading

Similar to the definition of model, the data_loader and val_data_loader for the training and validation stages are defined in the config configuration file using the data_loader and val_data_loader dictionaries, respectively, which correspond to the processing pipelines for the training and validation datasets.

Since QCNet does not include complex data augmentation (transforms=None), collate_fn defines how individual data items are grouped into batches, including alignment operations.

data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="Argoverse2PackedDataset",
        data_path=os.path.join(data_rootdir, "train_lmdb"),
        split="train",
        pack_type="lmdb",
        input_dim=2,
        transforms=None,
    ),
    sampler=dict(type=torch.utils.data.DistributedSampler),
    batch_size=batch_size_per_gpu,
    shuffle=True,
    num_workers=dataloader_workers,
    pin_memory=True,
    collate_fn=partial(
        collate_qc_argoverse2,
        ori_historical_sec=ori_historical_sec,
        ori_future_sec=ori_future_sec,
        ori_sample_fre=ori_sample_fre,
        sample_fre=sample_fre,
        stage="train",
        add_noise=False,
        agent_num=30,
        pl_N=80,
        pt_N=None,
        pt_N_downsample_nums=50,
    ),
)

val_data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="Argoverse2PackedDataset",
        data_path=os.path.join(data_rootdir, "val_lmdb"),
        split="val",
        pack_type="lmdb",
        input_dim=2,
        transforms=None,
    ),
    sampler=dict(type=torch.utils.data.DistributedSampler),
    batch_size=batch_size_per_gpu,
    shuffle=False,
    num_workers=dataloader_workers,
    pin_memory=True,
    collate_fn=partial(
        collate_qc_argoverse2,
        ori_historical_sec=ori_historical_sec,
        ori_future_sec=ori_future_sec,
        ori_sample_fre=ori_sample_fre,
        sample_fre=sample_fre,
        stage="val",
        add_noise=False,
        agent_num=30,
        pl_N=80,
        pt_N=50,
    ),
)

Additionally, the config defines batch_processor for processing batches:

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

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

In batch_processor, a loss_collector function is passed to collect the losses for the current batch of data, as shown below:

def loss_collector(outputs: dict):
    losses = []
    for _, loss in outputs.items():
        losses.append(loss)
    return losses   

Training Strategy

First, let's introduce the strategy for training the floating-point model of QCNet on the Argoverse2 dataset. We use AdamW as the optimizer with settings lr=5e-4 and weight_decay=1e-4. However, we do not apply weight decay to bias , nn.Embedding , and normalization layers. Therefore, using custom_param_optimizer here easily achieves this setup. We employ a Cosine learning rate schedule, set the warmup length to 1 epoch, and train the model for a total of 250 epochs. Below is a complete configuration example for the float_trainer in the config file:

float_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    data_loader=data_loader,
    optimizer=dict(
        type="custom_param_optimizer",
        optim_cls=torch.optim.AdamW,
        optim_cfgs=dict(lr=5e-4, weight_decay=1e-4),
        custom_param_mapper={
            "bias": dict(weight_decay=0.0),
            "norm_types": dict(weight_decay=0.0),
            nn.Embedding: dict(weight_decay=0.0),
        },
    ),
    batch_processor=batch_processor,
    num_epochs=250,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(
            type="CosLrUpdater",
            warmup_len=1,
            warmup_by="epoch",
            step_log_interval=500,
        ),
        val_callback,
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    sync_bn=True,
    val_metrics=val_metrics,
)

Quantized Model Training

For key steps in quantization training, such as preparing the floating-point model, operator replacement, inserting quantization and dequantization nodes, setting quantization parameters, and operator fusion, please refer to the Quantized Awareness Training (QAT) section. Here we focus on defining and using the quantized model.

The quantization training strategy for the QCNet example model can be found in the configs/traj_pred/qcnet_oe_argoverse2.py file, mainly divided into the quantization calibration calibration and quantized awareness training qat phases.

Quantization Calibration Configuration:

output_pred_L1_sensitive_ops = os.path.join(ckpt_dir, "analysis", "output_pred_L1_sensitive_ops.pt")
output_prob_L1_sensitive_ops = os.path.join(ckpt_dir, "analysis", "output_prob_L1_sensitive_ops.pt")

cali_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(
        observer=(
            observer_v2.MSEObserver
        )
    ),
    templates=[
        ModuleNameTemplate(
            {"": {"output": qint8, "weight": qint8}},
            enable_propagate=True,
        ),
        ModuleNameTemplate(
            {
                "backbone.quant": {"dtype": qint8, "threshold": 1.0},
                "side_backbone.quant": {"dtype": qint8, "threshold": 1.0},
                "tfl_bind_task_head.head.fusion_bev_feat_quant": {
                    "dtype": torch.float16
                },
            },
            freeze=True,
        ),
        MatmulDtypeTemplate(  # gemm single int8 input
            input_dtypes=[qint8, qint8],
        ),
        ConvDtypeTemplate(  # gemm int8 input
            input_dtype=qint8,
            weight_dtype=qint8,
        ),
        SensitivityTemplate(
            sensitive_table=output_pred_L1_sensitive_ops,
            topk_or_ratio=set_custom["output_pred_L1_sensitive_ops"],
        ),
        SensitivityTemplate(
            sensitive_table=output_prob_L1_sensitive_ops,
            topk_or_ratio=set_custom["output_prob_L1_sensitive_ops"],
        ),
        ModuleNameTemplate(
            {
                "decoder.loc_cumsum_conv": {"output": qint16},
                "encoder.agent_encoder.x_a_his_quant": {"output": qint16},
                "encoder.agent_encoder.x_a_emb.freqs.0": {
                    "dtype": {"input": qint16, "output": qint16}
                },
                "encoder.agent_encoder.x_a_emb.freqs.2": {
                    "dtype": {"input": qint16, "output": qint16}
                },
                "encoder.agent_encoder.x_a_emb._generated_mul_0": {
                    "output": qint16
                },
                "encoder.agent_encoder.x_a_emb._generated_mul_2": {
                    "output": qint16
                },
                "encoder.agent_encoder.r_t_quant.0": {"output": qint16},
                "encoder.agent_encoder.r_t_quant.3": {"output": qint16},
                "encoder.agent_encoder.x_a_quant.0": {"output": qint16},
                "decoder.locs_propose_pos_cat": {"output": qint16},
                "encoder.agent_encoder.mid_his_quant.0": {"output": qint16},
                "encoder.agent_encoder.mid_his_quant.0": {"threshold": 146.0},
                "encoder.agent_encoder.x_a_emb.to_out.2": {"output": qint16},
                "encoder.agent_encoder.x_a_emb.to_out.2": {"threshold": 146.0},
                
            },
            freeze=True,
            enable_propagate=True,
        ),
    ],
    enable_optimize=True,
    save_dir="./qconfig_setting",
    custom_qconfig_mapping={
        "encoder.agent_encoder.x_a_quant.2": get_default_qconfig(
            activation_observer="fixed_scale",
            activation_qkwargs={
                "dtype": qint16,
                "scale": 5 / 32768,
            },
        ),
        "encoder.agent_encoder.x_a_emb.freqs.2": get_default_qconfig(
            activation_observer="fixed_scale",
            activation_qkwargs={
                "dtype": qint16,
                "scale": 1.0 / 32768,
            },
        ),
    }.update(custom_qconfig_sincos),
)
float2calibration = dict(
    type="Float2Calibration",
    convert_mode="jit-strip",
    example_data_loader=calibration_example_data_loader,
    qconfig_setter=cali_qconfig_setter,
)

calibration_trainer = dict(
    type="Calibrator",
    model=cali_model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode=qat_mode,
        qconfig_params=dict(
            activation_calibration_observer="mse",
        ),
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "float-checkpoint-last.pth.tar"
                ),
                allow_miss=True,
                ignore_extra=True,
                verbose=True,
            ),
            float2calibration,
        ],
    ),
    data_loader=calibration_data_loader,
    batch_processor=calibration_batch_processor,
    num_steps=calibration_step,
    device=None,
    callbacks=[
        ckpt_callback,
        val_callback,
    ],
    log_interval=calibration_step / 10,
    val_metrics=val_metrics,
)

Quantized Awareness Training Configuration:

float2qat = dict(
    type="Float2QAT",
    convert_mode="jit-strip",
    example_data_loader=calibration_example_data_loader,
    qconfig_setter=qat_qconfig_setter,
)
qat_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=cali_model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode=qat_mode,
        converters=[
            float2qat,
            dict(
                type="LoadCheckpoint",
                checkpoint_path=os.path.join(
                    ckpt_dir, "calibration-checkpoint-last.pth.tar"
                ),
                verbose=True,
                allow_miss=True,
                ignore_extra=True,
            ),
        ],
    ),
    data_loader=qat_data_loader,
    optimizer=dict(
        type="custom_param_optimizer",
        optim_cls=torch.optim.AdamW,
        optim_cfgs=dict(lr=8 * 5e-6, weight_decay=1e-4),  # gpu: 8*4090
        custom_param_mapper={
            "bias": dict(weight_decay=0.0),
            "norm_types": dict(weight_decay=0.0),
            nn.Embedding: dict(weight_decay=0.0),
        },
    ),
    batch_processor=batch_processor,
    num_epochs=15,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(
            type="StepDecayLrUpdater",
            lr_decay_id=[4],
            step_log_interval=500,
        ),
        val_callback,
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=val_metrics,
)

Here, float2calibration and float2qat respectively define the conversion process from floating-point to calibration model and from floating-point to quantized awareness training model. The cali_qconfig_setter and qat_qconfig_setter provide the qconfig settings for the calibration model and the qat model respectively. For more details on setting and debugging qconfig, such as default qconfig settings and quantization-sensitive operator settings, please refer to the Quantized Awareness Training - Qconfig Details section.

Quantization Sensitivity Operator Sorting

We provide "output_prob_L1_sensitive_ops.pt" and "output_pred_L1_sensitive_ops.pt", and you can follow these steps to generate the sensitive tables:

First, you should set a calibration model with the default qconfig:

cali_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(
        observer=(
            observer_v2.MSEObserver
        )
    ),
    templates=[
        ModuleNameTemplate(
            {"": {"output": qint8, "weight": qint8}},
            enable_propagate=True,
        ),
        MatmulDtypeTemplate(
            input_dtypes=[qint8, qint8],
        ),
        ConvDtypeTemplate(
            input_dtype=qint8,
            weight_dtype=qint8,
        ),
    ],
    enable_optimize=True,
    save_dir="./qconfig_setting",
)

Then you should rename the model to "calibration-checkpoint-last.pth.tar". And calling the sensitivity analysis tool to compare the similarity between the float model and the calibration model layer by layer with the command:

python3 tools/quant_analysis.py --config configs/traj_pred/qcnet_oe_argoverse2.py

This command corresponds to following command:

analysis_convert_pipeline = dict(
    type="ModelConvertPipeline",
    qat_mode="fuse_bn",
    converters=[
        dict(
            type="Float2QAT",
            convert_mode="jit-strip",
            # example_data_loader=calibration_example_data_loader,
            qconfig_setter=cali_qconfig_setter,
            example_inputs=example_tmp_inputs,
        ),
        dict(
            type="LoadCheckpoint",
            checkpoint_path=os.path.join(
                ckpt_dir, "calibration-checkpoint-last.pth.tar"
            ),
            allow_miss=True,
            ignore_extra=True,
            verbose=True,
        ),
    ],
)
quant_data_loader = copy.deepcopy(val_data_loader)
quant_data_loader["batch_size"] = 1
quant_data_loader["num_workers"] = 0

quant_analysis_solver = dict(
    type="QuantAnalysis",
    model=cali_model,
    device_id=0,
    dataloader=quant_data_loader,
    num_steps=1000,
    baseline_model_convert_pipeline=float_predictor["model_convert_pipeline"],
    analysis_model_convert_pipeline=analysis_convert_pipeline,
    analysis_model_type="fake_quant",
    out_dir=os.path.join(ckpt_dir, "analysis"),
)

For the key steps, please refer to the Tools Guide - Accuracy Tuning Tool Guide section.