Bev Sparse Multitask Model Training

This tutorial mainly tells everyone how to train a BevSparseMultiTask model from scratch on the dataset nuscenes using HAT, including floating-point, quantized, and fixed-point models. This model is a bev multi-task perception model, which contains detection heads for three major perception tasks: sparse dynamic detection head, sparse static element detection head, and general obstacle detection head.

Training Process

If you just want to simply train the BevSparseMultiTask model, you can first read the content of this chapter. Similar to other tasks, for all training and evaluation tasks, HAT uniformly adopts the tools + config form to complete. After preparing the original dataset, you can conveniently complete the entire training process through the following process.

Dataset Preparation

For this example, using the nuScenes dataset, you can download the dataset from https://www.nuscenes.org/nuscenes. For the Occupancy prediction task, you will also need to download the OCC ground truth data, which can be obtained from https://github.com/CVPR2023-3D-Occupancy-Prediction/CVPR2023-3D-Occupancy-Prediction. It is recommended to unzip the downloaded dataset into the occ3d/gts directory within the nuScenes dataset folder. Additionally, to improve training speed, we have packaged the original JPEG format dataset into an LMDB format dataset.

python3 tools/dataset_converters/create_data_flashocc.py --src-data-path ./tmp_data/nuscenes/ -o ./tmp_data/nuscenes/occ3d

The previous command will generate two files, nuscenes_infos_train.pkl and nuscenes_infos_val.pkl, in the "./tmp_data/nuscenes/occ3d directory". Next, execute the following command to perform the packaging:

python3 tools/datasets/nuscenes_packer.py --src-data-dir WORKSAPCE/datasets/nuscenes/ --pack-type lmdb --target-data-dir . --version v1.0-trainval --split-name val --need-occ
python3 tools/datasets/nuscenes_packer.py --src-data-dir WORKSAPCE/datasets/nuscenes/ --pack-type lmdb --target-data-dir . --version v1.0-trainval --split-name train --need-occ

The above two commands correspond to converting the training dataset and the validation dataset, respectively. After the packaging is complete, the file structure in the data directory should be as follows:

tmp_data
    |-- nuscenes
        |-- v1.0-trainval 
            |-- train_lmdb
            |-- val_lmdb

train_lmdb and val_lmdb are the training dataset and validation dataset after packaging, which are also the final dataset read by the network. The meta contains initialization information needed by the evaluation script, with specific information copied from the original nuscenes dataset.

Model Training

After the dataset is prepared, you can start training the floating-point BevSparseMultiTask network. Before starting network training, you can use the following command to test the computation amount and parameter count of the network:

python3 tools/calops.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py --method hook

If you just want to start such a training task, you only need to run the following command:

python3 tools/train.py --stage "float" --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

Due to the clever registration mechanism used by the HAT algorithm package, each training task can be started in the form of train.py plus a config configuration file. train.py is a unified training script, task-independent. What kind of task we want to train, what dataset to use, and training-related hyperparameter settings are all contained in the specified config configuration file. The config file provides key dicts for model construction, data reading, etc.

Export Fixed-Point Model

After completing quantized training, you can start exporting the fixed-point model. This can be done through the following command:

python3 tools/export_hbir.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

Model Validation

After the model training is completed, we can also verify the performance of the trained model. Since we provide a two-stage training process of float, calibration, qat, accordingly we can verify the model performance trained in these three stages. Just run the following three commands respectively:

python3 tools/predict.py --stage "float" --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py
python3 tools/predict.py --stage "calibration" --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

Meanwhile, we also provide performance testing for quantized models. Just run the following command, but note that you must first export hbir:

python3 tools/predict.py --stage "int_infer" --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

The accuracy displayed here is the actual accuracy of the final quantized model. Of course, this accuracy should be very close to the accuracy in the qat validation phase.

Model Inference and Result Visualization

If you hope to see the visualization effect of the trained model, the tools folder below also provides a visualization script. You just need to run the following script:

python3 tools/infer_hbir.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py  --model-inputs ${model_inputs} --save-path ${save_path} --use-dataset

python3 tools/infer_hbir.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py  --model-inputs ${model_inputs} --save-path ${save_path}

Fixed-Point Model Inspection and Compilation

The quantization training toolchain integrated in HAT is mainly prepared for Horizon's computing platform. Therefore, the inspection and compilation of quantized models are necessary. We provide a model inspection interface in HAT, which can check whether it can run normally on the BPU after defining the quantized model:

python3 tools/model_checker.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

After training is completed, the compile_perf_hbir tool can be used to compile the quantized model into an hbm file that can run on the board. At the same time, this tool can also estimate the running performance on the BPU. The following script can be used:

python3 tools/compile_perf_hbir.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

Training Details

In this instruction, we explain some matters that need attention in model training, mainly some related settings of config.

Model Construction

The BevSparseMultiTask model contains three task heads: dynamic, static, and general obstacle. The dynamic task uses SparseBEVOEHead.


det_head = dict(
    type="SparseBEVOEHead",
    enable_dn=True,
    level_index=[2],
    cls_threshold_to_reg=0.05,
    instance_bank=dict(
        type="MemoryBankOE",
        num_anchor=384,
        embed_dims=embed_dims,
        num_memory_instances=384,
        anchor=anchor_file,
        num_temp_instances=128,
        confidence_decay=0.6,
    ),
    anchor_encoder=dict(
        type="SparseBEVOEEncoder",
        pos_embed_dims=128,
        size_embed_dims=32,
        yaw_embed_dims=32,
        vel_embed_dims=64,
        vel_dims=3,
    ),
    num_single_frame_decoder=num_single_frame_decoder,
    operation_order=[
        "deformable",
        "ffn",
        "norm",
        "refine",
    ]
    * num_single_frame_decoder
    + [
        "temp_interaction",
        "interaction",
        "norm",
        "deformable",
        "ffn",
        "norm",
        "refine",
    ]
    * (num_decoder - num_single_frame_decoder),
    ffn=dict(
        type="AsymmetricFFNOE",
        in_channels=embed_dims * 2,
        pre_norm=True,
        embed_dims=embed_dims,
        feedforward_channels=embed_dims * 4,
        num_fcs=2,
        ffn_drop=drop_out,
    ),
    deformable_model=dict(
        type="DeformableFeatureAggregationOE",
        embed_dims=embed_dims,
        num_groups=num_groups,
        num_levels=num_levels,
        num_cams=6,
        attn_drop=0.15,
        use_camera_embed=True,
        residual_mode="cat",
        kps_generator=dict(
            type="SparseBEVOEKeyPointsGenerator",
            num_pts=8,
        ),
    ),
    refine_layer=dict(
        type="SparseBEVOERefinementModule",
        embed_dims=embed_dims,
        num_cls=num_classes,
        refine_yaw=True,
    ),
    target=dict(
        type="SparseBEVOETarget",
        num_dn_groups=5,
        num_temp_dn_groups=3,
        dn_noise_scale=[2.0] * 3 + [0.5] * 7,
        max_dn_gt=32,
        add_neg_dn=True,
        cls_weight=2.0,
        box_weight=0.25,
        reg_weights=[2.0] * 3 + [0.5] * 3 + [0.0] * 4,
        cls_wise_reg_weights={
            CLASSES.index("traffic_cone"): [
                2.0,
                2.0,
                2.0,
                1.0,
                1.0,
                1.0,
                0.0,
                0.0,
                1.0,
                1.0,
            ],
        },
    ),
    cls_allow_reverse=[CLASSES.index("barrier")],
    loss_cls=dict(
        type="FocalLoss",
        loss_name="cls",
        num_classes=num_classes + 1,
        gamma=2.0,
        alpha=0.25,
        loss_weight=2.0,
    ),
    loss_reg=dict(type="L1Loss", loss_weight=0.25),
    loss_cns=dict(type="CrossEntropyLoss", use_sigmoid=True),
    loss_yns=dict(type="GaussianFocalLoss"),
    decoder=dict(type="SparseBEVOEDecoder"),
    reg_weights=[2.0] * 3 + [1.0] * 7,
)

The static task uses SparseMapPerceptionDecoder.


om_head = dict(
    type="SparseMapPerceptionDecoder",
    embed_dims=embed_dims,
    num_cam=6,
    num_vec_one2one=om_num_anchor,
    num_vec_one2many=0,
    k_one2many=0,
    num_pts_per_vec=fixed_ptsnum_per_pred_line,
    transform_method="minmax",
    is_deploy=False,
    decoder=dict(
        type="SparseOMOEHead",
        feat_indices=[2],
        num_anchor=om_num_anchor,
        num_pts_per_vec=fixed_ptsnum_per_pred_line,
        num_views=6,
        projection_mat_key="lidar2img",
        anchor_encoder=dict(
            type="SparseOEPoint3DEncoder",
            embed_dims=embed_dims,
            input_dim=fixed_ptsnum_per_pred_line * 2,
        ),
        instance_bank=dict(
            type="InstanceBankOE",
            num_anchor=om_num_anchor,
            embed_dims=embed_dims,
            anchor=map_anchor_file,
            num_temp_instances=0,
            confidence_decay=0.6,
            anchor_grad=True,
            feat_grad=True,
        ),
        instance_interaction=dict(
            type=MultiheadAttention,
            embed_dim=embed_dims,
            num_heads=num_groups,
            batch_first=True,
            attn_drop=drop_out,
            proj_drop=drop_out,
        ),
        norm_layer=dict(
            type=torch.nn.LayerNorm,
            normalized_shape=embed_dims,
        ),
        ffn=dict(
            type="AsymmetricFFNOE",
            in_channels=embed_dims * 2,
            pre_norm=True,
            embed_dims=embed_dims,
            num_fcs=2,
            ffn_drop=0.1,
            feedforward_channels=embed_dims * 4,
        ),
        deformable_model=dict(
            type="DeformableFeatureAggregationOEv2",
            embed_dims=embed_dims,
            num_groups=num_groups,
            num_levels=num_levels,
            num_cams=6,
            attn_drop=0.15,
            residual_mode="cat",
            use_camera_embed=True,
            grid_align_num=4,
            kps_generator=dict(
                type="SparsePoint3DKeyPointsGenerator",
                embed_dims=embed_dims,
                num_sample=fixed_ptsnum_per_pred_line,
                num_learnable_pts=2,
                fix_height=(0, 0.5, -0.5),
                ground_height=-1.84023,  # ground height in lidar frame
            ),
        ),
        refine_layer=dict(
            type="SparsePoint3DRefinementModule",
            embed_dims=embed_dims,
            num_sample=fixed_ptsnum_per_pred_line,
            coords_dim=2,
            num_cls=len(map_classes),
            with_cls_branch=True,
        ),
        num_decoder=num_decoder,
        num_single_frame_decoder=num_single_frame_decoder,
        operation_order=[
            "interaction",
            "norm",
            "deformable",
            "ffn",
            "norm",
            "refine",
        ]
        * num_single_frame_decoder
        + [
            "interaction",
            "norm",
            "deformable",
            "ffn",
            "norm",
            "refine",
        ]
        * (num_decoder - num_single_frame_decoder),
    ),
    aux_seg=map_aux_seg_cfg,
    criterion=dict(
        type="MapTRCriterion",
        dir_interval=1,
        num_classes=num_map_classes,
        code_weights=[1.0, 1.0, 1.0, 1.0],
        sync_cls_avg_factor=True,
        pc_range=map_point_cloud_range,
        num_pts_per_vec=fixed_ptsnum_per_pred_line,  # one bbox
        num_pts_per_gt_vec=fixed_ptsnum_per_gt_line,
        gt_shift_pts_pattern="sparsedrive",
        aux_seg=map_aux_seg_cfg,
        pred_absolute_points=True,
        assigner=dict(
            type="MapTRAssigner",
            cls_cost=dict(type="FocalLossCost", weight=1.0),
            pts_cost=dict(type="OrderedPtsL1Cost", weight=0.5, beta=0.01),
            pc_range=map_point_cloud_range,
            pred_absolute_points=True,
        ),
        loss_cls=dict(
            type="FocalLoss",
            loss_name="cls",
            num_classes=num_map_classes + 1,
            alpha=0.25,
            gamma=2.0,
            loss_weight=1.0,
            reduction="mean",
        ),
        loss_pts=dict(type="PtsL1Loss", loss_weight=0.5, beta=0.01),
        loss_dir=dict(type="PtsDirCosLoss", loss_weight=0),
        loss_pv_seg=dict(type="SimpleLoss", pos_weight=1.0, loss_weight=2.0),
    ),
    post_process=dict(
        type="MapTRPostProcess",
        pc_range=map_point_cloud_range,
        post_center_range=map_post_center_range,
        max_num=om_num_anchor,
        num_classes=num_map_classes,
        pred_absolute_points=True,
    ),
)

The head selected for the general obstacle task is FlashOccHead.

occ_head = dict(
    type="FlashOccHead",
    bev_feat_index=-1,
    bev_upscale=2,
    view_transformer=dict(
        type="LSSTransformer",
        in_channels=256,
        feat_channels=64,
        z_range=(-1.0, 5.4),
        depth=occ_depth,
        num_points=occ_num_points,
        bev_size=occ_bev_size,
        grid_size=occ_grid_size,
        num_views=6,
        grid_quant_scale=grid_quant_scale,
        depth_grid_quant_scale=depth_grid_quant_scale,
        homo_key="lidar2img",
    ),
    bev_encoder=dict(
        type="BevEncoder",
        backbone=dict(
            type="HENet",
            in_channels=64,
            block_nums=[4, 3, 8, 6],
            embed_dims=[64, 128, 192, 384],
            attention_block_num=[0, 0, 0, 0],
            mlp_ratios=[2, 2, 2, 3],
            mlp_ratio_attn=2,
            act_layer=["nn.GELU", "nn.GELU", "nn.GELU", "nn.GELU"],
            use_layer_scale=[True, True, True, True],
            layer_scale_init_value=1e-5,
            num_classes=1000,
            include_top=False,
            extra_act=[False, False, False, False],
            final_expand_channel=0,
            feature_mix_channel=1024,
            block_cls=["GroupDWCB", "GroupDWCB", "AltDWCB", "DWCB"],
            down_cls=["S2DDown", "S2DDown", "S2DDown", "None"],
            patch_embed="origin",
            quant_input=False,
        ),
        neck=dict(
            type="BiFPN",
            in_strides=[2, 4, 8, 16, 32],
            out_strides=[2, 4, 8, 16, 32],
            stride2channels=dict({2: 64, 4: 64, 8: 128, 16: 192, 32: 384}),
            out_channels=48,
            num_outs=5,
            stack=3,
            start_level=0,
            end_level=-1,
            fpn_name="bifpn_sum",
            upsample_type="function",
            use_fx=False,
        ),
    ),
    bev_decoder=dict(
        type="FlashOccDetDecoder",
        use_mask=True,
        num_classes=num_classes_occ,
        occ_head=dict(
            type="BEVOCCHead2D",
            in_dim=48,
            out_dim=128,
            Dz=16,
            num_classes=num_classes_occ,
            use_predicter=True,
            use_upsample=True,
        ),
        loss_occ=dict(
            type="CrossEntropyLoss",
            use_sigmoid=False,
            ignore_index=255,
            loss_weight=6.0,
        ),
    ),
)

The heads of the three tasks together with the backbone constitute this BEV multi-task perception model. The model is defined in the config as:


task_heads = OrderedDict()
if enable_det_head:
    task_heads["det"] = det_head
if enable_om_head:
    task_heads["om"] = om_head
if enable_occ_head:
    task_heads["occ"] = occ_head

model = dict(
    type="SparseBevFusionMultitaskOE",
    compiler_model=False,
    backbone=dict(
        type="HENet",
        in_channels=3,
        block_nums=[4, 3, 8, 6],
        embed_dims=[64, 128, 192, 384],
        attention_block_num=[0, 0, 0, 0],
        mlp_ratios=[2, 2, 2, 3],
        mlp_ratio_attn=2,
        act_layer=["nn.GELU", "nn.GELU", "nn.GELU", "nn.GELU"],
        use_layer_scale=[True, True, True, True],
        layer_scale_init_value=1e-5,
        num_classes=1000,
        include_top=False,
        extra_act=[False, False, False, False],
        final_expand_channel=0,
        feature_mix_channel=1024,
        block_cls=["GroupDWCB", "GroupDWCB", "AltDWCB", "DWCB"],
        down_cls=["S2DDown", "S2DDown", "S2DDown", "None"],
        patch_embed="origin",
    ),
    neck=dict(
        type="MMFPN",
        in_strides=[2, 4, 8, 16, 32],
        in_channels=[64, 64, 128, 192, 384],
        fix_out_channel=256,
        out_strides=[4, 8, 16, 32],
    ),
    depth_branch=dict(  # for auxiliary supervision only
        type="DenseDepthNetOE",
        embed_dims=embed_dims,
        num_depth_layers=num_depth_layers,
        loss_weight=0.2,
    ),
    head=task_heads,
)

Among them, the type under model represents the defined model name, and the remaining variables represent other components of the model. The advantage of defining a model in this way is that we can easily replace the structure we want. After the training script starts, it will call the build_model interface to turn such a dict-type model into a model of type torch.nn.Module.

Data Augmentation

Similar to the definition of model, the data augmentation process is implemented by defining data_loader and val_data_loader these two dicts in the config configuration file, corresponding to the processing procedures of training set and validation set respectively. The Nuscenes dataset simultaneously contains annotations for dynamic, static, and general obstacles. We can train three tasks on one dataset. In the specific code implementation, NuscenesBevDataset can provide dynamic and occ annotations, and NuscenesSparseMapDataset inherits from it and can provide dynamic, static, and occ annotations simultaneously. The dynamic detection head uses temporal information, so data_loader adopts DistStreamBatchSampler to ensure that data is returned in sequence.

train_dataset = dict(
    type="NuscenesSparseMapDataset",
    data_path=os.path.join(data_rootdir, "train_lmdb"),
    map_path=meta_rootdir,
    pc_range=map_point_cloud_range,
    test_mode=False,
    bev_size=(bev_h_, bev_w_),
    fixed_ptsnum_per_line=fixed_ptsnum_per_gt_line,
    padding_value=-10000,
    map_classes=map_classes,
    aux_seg=map_aux_seg_cfg,
    with_bev_bboxes=False,
    with_ego_bboxes=False,
    with_bev_mask=False,
    with_lidar_bboxes=True,
    need_lidar=True,
    num_split=2,
    num_sweeps=9,
    load_dim=5,
    use_dim=[0, 1, 2, 3, 4],
    use_lidar_gt=True,
    with_ego_occ=False,
    with_lidar_occ=True,
    filter_empty_gt=False,
    transforms=[
        dict(type="MultiViewsImgResize", scales=(0.40, 0.47)),
        dict(type="MultiViewsImgCrop", size=(256, 704), random=False),
        dict(type="MultiViewsImgFlip"),
        dict(type="MultiViewsImgRotate", rot=(-5.4, 5.4)),
        dict(type="BevBBoxRotation", rotation_3d_range=(-0.3925, 0.3925)),
        dict(type="MultiViewsPhotoMetricDistortion", use_pil=False),
        dict(
            type="MultiViewsGridMask",
            use_h=True,
            use_w=True,
            rotate=1,
            offset=False,
            ratio=0.5,
            mode=1,
            prob=0.7,
        ),
        dict(
            type="MultiViewsImgTransformWrapper",
            transforms=[
                dict(type="PILToTensor"),
                dict(type="BgrToYuv444", rgb_input=True),
                dict(type="Normalize", mean=128, std=128),
            ],
        ),
    ],
)

data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=train_dataset,
    batch_sampler=dict(
        type="DistStreamBatchSampler",
        batch_size=batch_size_per_gpu,
        dataset=train_dataset,
        keep_consistent_seq_aug=True,
        skip_prob=0.0,
        sequence_flip_prob=0.0,
    ),
    num_workers=dataloader_workers,
    pin_memory=True,
    collate_fn=collate_nuscenes,
)

Training Strategy

To train a model with high accuracy, a good training strategy is indispensable. For each training task, the corresponding training strategy is also defined in the config file, which can be seen from the variable 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=os.path.join(
                    pretrain_ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
                allow_miss=True,
                ignore_extra=True,
                ignore_tensor_shape=True,
                verbose=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        params={
            "backbone": dict(lr=3e-4),
        },
        # eps=1e-8,
        # betas=(0.9, 0.999),
        lr=3e-4,
        weight_decay=0.001,
    ),
    batch_processor=batch_processor,
    num_steps=num_steps,
    stop_by="step",
    # num_epochs=100,
    device=None,
    callbacks=[
        stat_callback,
        loss_show_update,
        dict(type="ExponentialMovingAverage"),
        grad_callback,
        dict(
            type="CosineAnnealingLrUpdater",
            warmup_len=500,
            warmup_by="step",
            warmup_lr_ratio=1.0 / 3,
            step_log_interval=500,
            update_by="step",
            min_lr_ratio=1e-3,
        ),
        val_callback,
        ckpt_callback,
    ],
    train_metrics=dict(
        type="LossShow",
    ),
    sync_bn=True,
    val_metrics=val_metrics,
)

float_trainer defines our training method on a large scale, including using multi-GPU distributed training (distributed_data_parallel_trainer), the number of epochs for model training, and the choice of optimizer. At the same time, callbacks reflects the small strategies used in the training process and operations you want to implement, including exponential moving average training (ExponentialMovingAverage), learning rate transformation method (CosineAnnealingLrUpdater), Validating model metrics during the training process (Validation), and saving (Checkpoint) model operations. Of course, if you have operations you hope the model to implement during the training process, you can also add them in the form of this dict. float_trainer is responsible for linking the entire training logic together.

Note

If you need to reproduce the accuracy, it is best not to modify the training strategy in the config. Otherwise, unexpected training situations may occur.

Through the above introduction, you should have a relatively clear understanding of the functions of the config file. Then, through the training script mentioned earlier, you can train a high-accuracy pure floating-point detection model. Of course, training a good detection model is not our ultimate goal; it just serves as a pretrain to serve our later training of fixed-point models.

Quantized Model Training

After we have a pure floating-point model, we can start quantized model training. Similar to floating-point training, we can use the following command to calibrate the model and obtain a pseudo-quantized model.

python3 tools/train.py --stage "calibration" --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

It can be seen that our configuration file has not changed, only the stage type has changed. At this time, the training strategy we use comes from calibration_trainer in the config file.


calibration_trainer = dict(
    type="Calibrator",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        qat_mode="fuse_bn",
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path=float_checkpoint,
                allow_miss=True,
                ignore_extra=True,
                verbose=True,
                load_ema_model=True,
            ),
            dict(
                type="Float2Calibration",
                convert_mode=convert_mode,
                example_inputs=example_inputs,
                qconfig_setter=calibration_qconfig_setter,
            ),
        ],
    ),
    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,
)

When we train the quantized model, we need to set model_convert_pipeline. It first loads the model's floating-point weights and then uses Float2Calibration to convert the model from a floating-point model to a calibration model. Among them, calibration_qconfig_setter defines the quantization configuration of the model. The current configuration is:


sensitivity_templates = []
table_dir = os.path.join(ckpt_dir, "quant_analysis")
sensitive_setting_dict = {
    "output_det.classification_L1_sensitive_ops.pt": 20,
    "output_det.prediction_L1_sensitive_ops.pt": 20,
    "output_det.quality_L1_sensitive_ops.pt": 10,
    "output_om.all_pts_preds_L1_sensitive_ops.pt": 5,
    "output_om.all_bbox_preds_L1_sensitive_ops.pt": 5,
    "output_om.all_cls_scores_L1_sensitive_ops.pt": 5,
}
for k, v in sensitive_setting_dict.items():
    sensitive_table = load_ckpt(os.path.join(table_dir, k))
    sensitivity_templates.append(
        SensitivityTemplate(
            sensitive_table,
            topk_or_ratio=v,
        )
    )

backbone = ["backbone"]
int16_max_point_muls = [
    f"head.det.layers.{layer}.point_mul" for layer in range(3, 39, 7)
] + [f"head.om.decoder.layers.{layer}.point_mul" for layer in range(2, 33, 6)]
int16_max_reciprocal_op = [
    f"head.det.layers.{layer}.reciprocal_op" for layer in range(3, 39, 7)
]
int16_add_ops = [f"head.det.layers.{layer}.add2" for layer in range(6, 35, 7)]
# int16_add_ops += [f"head.om.layers.{layer}._generated_add_1" for layer in range(5, 35, 6)]
float32_ops = [
    f"head.det.layers.{layer}.cls_layers" for layer in range(6, 35, 7) # 6, 13, 20, 27, 34
] + [f"head.det.layers.{layer}.quality_layers" for layer in range(6, 35, 7)]
q_template = [
    ModuleNameTemplate({"": qint8}),
    MatmulDtypeTemplate(
        input_dtypes=[qint8, qint8],
        # prefix=[""],
    ),
    ConvDtypeTemplate(
        input_dtype=qint8,
        weight_dtype=qint8,
        # prefix=[""],
    ),
    ModuleNameTemplate(
        {
            m: {"dtype": qint8, "threshold": 1.0} for m in ["backbone.quant"]
        },

    ),
    ModuleNameTemplate(
        {
            m: {"dtype": qint16, "threshold": 1.1}
            for m in int16_max_point_muls
        },

    ),
    ModuleNameTemplate(
        {
            m: {"dtype": qint16, "threshold": 10}
            for m in int16_max_reciprocal_op
        },

    ),
    ModuleNameTemplate(
        {
            m: {"dtype": qint16, "threshold": 50}
            for m in ["head.det.fc_after"]
        },

    ),
    ModuleNameTemplate(
        {
            m: {"dtype": qint16, "threshold": 10}
            for m in ["head.det.fc_before"]
        },
    ),
    ModuleNameTemplate(
        {
            m: {"input": [qint16, qint16], "output": qint16}
            for m in int16_add_ops
        },

    ),
    *sensitivity_templates,
    ModuleNameTemplate(
        {m: None for m in float32_ops},  # quant int8, fixed scale method 1
        freeze=True,
    ),
]

calibration_qconfig_setter = QconfigSetter(
    reference_qconfig=get_qconfig(  # 1. Mainly used to obtain observer
        # observer=(observer_v2.MSEObserver)
        observer=(observer_v2.HistogramObserver)
    ),
    templates=q_template,
    enable_optimize=True,
    save_dir=os.path.join(
        ckpt_dir, "./qconfig_setting"
    ),
    custom_qconfig_mapping=None,
)

Explanation of Quantization Configuration

The quantization configuration here is slightly complex, and we will explain it in detail step by step.

First, the first parameter of calibration_qconfig_setter reference_qconfig defines a basic quantization configuration, using the observer observer_v2.HistogramObserver. This is a calibration tool based on histogram statistics. It globally statistically analyzes the histogram statistics of the input and output features of each layer of the model, and then determines the quantization step size scale for each layer.

The second parameter of calibration_qconfig_setter is templates corresponding to q_template, which uses many templates to define the quantization configuration from top to bottom. At the top is ModuleNameTemplate({"": qint8}), defining the global quantization configuration as qint8, then the two lines MatmulDtypeTemplate and ConvDtypeTemplate define the input types for matmul operators and conv operators as qint8 respectively. Then, many lines use ModuleNameTemplate to specially set the quantization type and fixed scale for many ops in the model, which is a configuration based on observations of the model structure and statistics. Users can also similarly modify the quantization configuration of each layer of the model. Note that there is a line *sensitivity_templates in q_template, which is a high-precision template configured based on sensitivity analysis results.

Although based on experience, we have already manually configured some high-precision operators, it is still difficult for quantized models to recover the accuracy of the floating-point model. To analyze the problems of model quantization, we suggest users perform sensitivity analysis on the model. To run the model sensitivity analysis tool, you can run the following command:

python3 tools/quant_analysis.py --config configs/bev/bev_sparse_det_maptr_flashocc_henet_tinym_nuscenes.py

This command corresponds to the following configuration in the config:

quant_analysis_solver = dict(
    type="QuantAnalysis",
    model=analysis_model,
    device_id=-1,
    dataloader=analysis_dataloader,
    num_steps=100,
    baseline_model_convert_pipeline=float_predictor["model_convert_pipeline"],
    analysis_model_convert_pipeline=calibration_predictor[
        "model_convert_pipeline"
    ],
    analysis_model_type="fake_quant",
    out_dir=os.path.join(ckpt_dir, "quant_analysis"),
)

The main principle of sensitivity analysis is: given data, the tool runs the floating-point model and quantized model separately, finds the samples with the largest output difference, known as bad case. Then for the quantized model, the tool traverses each operator in the model, performs single operator quantization, sorts based on the error comparison of the floating-point model output after quantization, obtaining a list of model quantization-sensitive operators.

sensitivity_templates precisely reads this sensitivity ranking list, configuring the operators ranked at the top of sensitivity to high precision, thereby improving the quantization accuracy of the model.

After calibration, the quantization accuracy of this model has reached over 99% of floating-point training. We can achieve the accuracy target without performing additional quantization-aware training.

For more information about model quantization, please read the Tools Guide - QAT Accuracy Tuning Tool section.