#Occupancy Prediction Model

The Occupancy prediction model is developed based on Horizon Torch Samples, a deep learning framework developed by Horizon. For information on using Horizon Torch Samples, please refer to the Horizon Torch Samples documentation. The training configuration for the Occupancy prediction model is located in the "configs/occ/" directory. The following example uses "configs/occ/flashocc_henet_lss_occ3d_nuscenes.py" to illustrate how to configure and train the Occupancy prediction model.

#Training Process

If you want to quickly train the flashocc_henet_lss_occ3d_nuscenes model, you can start by reading this section. Like other tasks, for all training and evaluation tasks, HAT uniformly uses the tools + config approach to complete them. After preparing the original dataset, you can easily follow the steps below to complete the entire training 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
        |-- metas
        |-- v1.0-trainval 
        |-- occ3d
            | -- gts
            | -- nuscenes_infos_train.pkl
            | -- nuscenes_infos_val.pkl
    |-- occ3d_nuscenes
        |-- train_lmdb
        |-- val_lmdb

"train_lmdb" and "val_lmdb" are the packaged training dataset and validation dataset, respectively. They are also the datasets that the network will ultimately read.

#Model Training

Once the dataset is prepared, you can start training the floating-point Occupancy prediction model.

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

python3 tools/train.py --stage "float" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py
python3 tools/train.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

The commands above separately train the floating-point model and the fixed-point (quantized) model. The training of the fixed-point model needs to be based on the previously trained floating-point model. For more details, please refer to the Quantized Awareness Training section.

#Export the Fixed-point Model

After completing quantization training, you can start exporting the quantized model. You can use the following command to export it:

python3 tools/export_hbir.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

#Model Validation

After training is complete, you will obtain floating-point, quantized, or fixed-point models. Similar to the training process, you can validate the trained models using the same method to obtain metrics for Float, Calibration, and Quantized, corresponding to the floating-point, quantized, and fully fixed-point metrics, respectively.

python3 tools/predict.py --stage "float" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

python3 tools/predict.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

Similar to when training the model, when the parameter after --stage is "float" or "calibration", you can validate the trained floating-point or quantized models, respectively.

You can also validate the accuracy of the fixed-point model using the command below, but it’s important to note that you must first export the HBIR:

python3 tools/predict.py --stage "int_infer" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

#Model Inference and Results Visualization

HAT provides the infer_hbir.py script to visualize the inference results of the quantized model:

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

python3 tools/infer_hbir.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py --model-inputs ${model_inputs} --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/occ/flashocc_henet_lss_occ3d_nuscenes.py

After model training is complete, the compile_perf_hbir script can compile the quantized model into an hbm file suitable for board deployment. This tool also estimates the model's performance on the BPU:

python3 tools/compile_perf_hbir.py --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

This outlines 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

model = dict(
    type="ViewFusion",
    bev_feat_index=-1,
    bev_upscale=2,
    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="FPN",
        in_strides=[2, 4, 8, 16, 32],
        in_channels=[64, 64, 128, 192, 384],
        out_strides=[16, 32],
        out_channels=[256, 256],
        bn_kwargs=dict(eps=1e-5, momentum=0.1),
    ),
    view_transformer=dict(
        type="LSSTransformer",
        in_channels=256,
        feat_channels=64,
        z_range=(-1.0, 5.4),
        depth=depth,
        num_points=num_points,
        bev_size=bev_size,
        grid_size=grid_size,
        num_views=6,
        grid_quant_scale=grid_quant_scale,
        depth_grid_quant_scale=depth_quant_scale,
    ),
    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=True,
        ),
    ),
    bev_decoders=[
        dict(
            type="FlashOccDetDecoder",
            use_mask=True,
            num_classes=num_classes,
            occ_head=dict(
                type="BEVOCCHead2D",
                in_dim=48,
                out_dim=128,
                Dz=16,
                num_classes=num_classes,
                use_predicter=True,
                use_upsample=True,
            ),
            loss_occ=dict(
                type="CrossEntropyLoss",
                use_sigmoid=False,
                ignore_index=255,
                loss_weight=1.0,
            ),
        ),
    ],
)

Here, the type under model specifies the name of the defined model, while the remaining variables represent other components of the model. The advantage of this model definition is that it allows us to easily replace the structure as needed. For example, if we want to train a model with a ResNet50 backbone, we only need to replace the backbone under model.

#Data Augmentation

Similar to the model definition, the data augmentation process is implemented by defining two dicts in the config file: data_loader and val_data_loader, which correspond to the processing workflows for the training and validation datasets, For detailed information, see "configs/occ/flashocc_henet_lss_occ3d_nuscenes.py".

#Training Strategy

A good training strategy is essential for training a high-accuracy model. For each training task, the corresponding training strategy is defined in the config file, as indicated by the float_trainer variable.

float_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    model_convert_pipeline=dict(
        type="ModelConvertPipeline",
        converters=[
            dict(
                type="LoadCheckpoint",
                checkpoint_path="./tmp_pretrained_models/henet_tinym_imagenet/float-checkpoint-best.pth.tar",
                allow_miss=True,
                ignore_extra=True,
                ignore_tensor_shape=True,
            ),
        ],
    ),
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.AdamW,
        lr=1e-4,
        weight_decay=1e-2,
    ),
    batch_processor=batch_processor,
    num_epochs=num_epochs,
    device=None,
    callbacks=[
        stat_callback,
        grad_callback,
        dict(
            type="CosineAnnealingLrUpdater",
            warmup_len=500,
            warmup_by="step",
            warmup_lr_ratio=1.0 / 3,
            step_log_interval=500,
            stop_lr=2e-4 * 1e-3,
        ),
        metric_updater,
        val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
    train_metrics=dict(
        type="LossShow",
    ),
    val_metrics=dict(
        type="MeanIOU",
        seg_class=occ3d_seg_class,
        ignore_index=17,
    ),
)

float_trainer defines our training approach from a broad perspective, including using multi-card distributed training (distributed_data_parallel_trainer), the number of epochs for model training, and the choice of optimizer. Meanwhile, callbacks reflect the smaller strategies and operations you wish to implement during training. This includes learning rate adjustment methods (e.g., CosineAnnealingLrUpdater), validating model metrics during training (Validation), and saving model checkpoints (Checkpoint). If you have additional operations you'd like to implement during training, you can also add them in this dict format.

From the above introduction, you should have a clear understanding of the config file's functionality. Using the training scripts mentioned earlier, you can train a high accuracy floating-point detection model. Of course, training a good detection model is not our final goal, it serves as a pretrain for training the fixed-point model later.

#Quantized Model Training

Once we have the floating-point model, we can start training the corresponding fixed-point model. Similar to floating-point training, you can obtain a pseudo-quantized model by running the following scripts, which will use calibration and qat to achieve the target:

python3 tools/train.py --stage "calibration" --config configs/occ/flashocc_henet_lss_occ3d_nuscenes.py

As you can see, our config file remains unchanged, only the stage type has been altered. The training strategy we use at this point 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=os.path.join(
                    ckpt_dir, "float-checkpoint-best.pth.tar"
                ),
            ),
            dict(type="Float2Calibration", convert_mode=convert_mode),
            dict(
                type="FixWeightQScale",
            ),
        ],
    ),
    data_loader=calibration_data_loader,
    batch_processor=calibration_batch_processor,
    num_steps=calibration_step,
    device=None,
    callbacks=[
        stat_callback,
        calibration_val_callback,
        calibration_ckpt_callback,
    ],
    val_metrics=dict(
        type="MeanIOU",
        seg_class=occ3d_seg_class,
        ignore_index=17,
    ),
    log_interval=calibration_step / 10,
)

#With Different quantize Parameter

When training a quantized model, you need to set quantize=True. At this point, the corresponding floating-point model will be converted into a quantized model. The relevant code is as follows:

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

For key steps in quantization training, e.g., preparing the floating-point model, operator substitution, inserting quantization and inverse quantization nodes, setting quantization parameters, and operator fusion, etc., please read the Quantized Awareness Training (QAT) section.

#With Different Training Strategies

As previously mentioned, the quantization training is actually the finetuning based on the pure floating-point training. Therefore, in the quantization training, set the initial learning rate is to one-tenth of the floating-point training, and the number of epochs of the training will greatly decrease as well. The most important thing is that when defining model, we need to set pretrained to the address of the trained pure floating-point model.

After these simple adjustments, we can start training our quantized model.