English

UNet Segmentation Model Training

As a sample of the HAT segmentation task, this tutorial demonstrates how to train a state-of-the-art floating-point and fixed-point model on the Cityscapes dataset using HAT.

Cityscapes is an image dataset of urban driving scenes, containing 5000 images with pixel-level annotations, and the objects in the images are divided into 19 categories. The segmentation task is relatively complex and requires a high level of model capability, and it is not easy to achieve good metrics on the segmentation task using small models.

This tutorial will elaborate on how to train a state-of-the-art segmentation model on the Cityscapes dataset using HAT and conduct quantitative training on top of the floating-point model to finally get a fixed-point model from scratch.

Training Process

Downloading Dataset

To download the Cityscapes dataset, you first need to register an account on the official website.

After that, you can download the dataset files from the Download page. Here we only need two files: gtFine_trainvaltest.zip and leftImg8bit_trainvaltest.zip.

Additionally, the official Cityscapes dataset provides a script for downloading and processing the data, see Github.

First install the official tools using the following command:

python3 -m pip install cityscapesscripts

Then use the official tool to log into the account registered above and download the required dataset files.

csDownload -d ${data-dir} gtFine_trainvaltest.zip
csDownload -d ${data-dir} leftImg8bit_trainvaltest.zip

Finally, unpack the downloaded files (optional):

cd ${data-dir}
unzip gtFine_trainvaltest.zip
unzip leftImg8bit_trainvaltest.zip

Packing Dataset

To improve data reading efficiency, we recommend pre-packing the dataset into the LMDB format. HAT provides the cityscapes_packer.py script for you to easily convert the dataset from the original public format to numpy.ndarray or torch.Tensor. Use msgpack to wrap the data and eventually pack it into LMDB files.

The command to pack the dataset is as follows:

python3 tools/datasets/cityscapes_packer.py --src-data-dir ${data-dir} --split-name train --pack-type lmdb
python3 tools/datasets/cityscapes_packer.py --src-data-dir ${data-dir} --split-name val --pack-type lmdb

The generated LMDB files are saved under ${data-dir}/train_lmdb and ${data-dir}/val_lmdb.

Model Training

After packing the dataset into LMDB files, you can start training the model. HAT provides the training script train.py to facilitate model training together with config files.

Before starting the training, please make sure that the dataset path (data_rootdir) in the configuration file unet_mobilenetv1_cityscapes.py is set to the LMDB file path.

The commands for model training are as follows:

python3 tools/train.py --stage float --config configs/segmentation/unet_mobilenetv1_cityscapes.py
python3 tools/train.py --stage calibration --config configs/segmentation/unet_mobilenetv1_cityscapes.py
python3 tools/train.py --stage qat --config configs/segmentation/unet_mobilenetv1_cityscapes.py

The above two commands are respectively for the training of the floating-point model and the fixed-point model. The fixed-point model training needs to be based on the trained fixed-point model. For details on this, please read the Quantized Awareness Training (QAT) section.

Metrics Validation

After training, the script will automatically validate the metrics of the trained model.

In addition, HAT provides a script to validate the metrics of the trained floating-point and fixed-point models, which is also configured by using the config file.

python3 tools/predict.py --stage float --config configs/segmentation/unet_mobilenetv1_cityscapes.py --ckpt ${float-checkpoint-path}
python3 tools/predict.py --stage calibration --config configs/segmentation/unet_mobilenetv1_cityscapes.py --ckpt ${calibration-checkpoint-path}
python3 tools/predict.py --stage qat --config configs/segmentation/unet_mobilenetv1_cityscapes.py --ckpt ${qat-checkpoint-path}
python3 tools/predict.py --stage int_infer --config configs/segmentation/unet_mobilenetv1_cityscapes.py --ckpt ${int-infer-checkpoint-path}

Simulation of On-board Accuracy Validation

In addition to the above model validation, we also provide the exact same accuracy validation method simulating the on-board conditions, as below:

python3 tools/align_bpu_validation.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Model Inference

HAT provides the infer.py script to visualize the inference results of the trained models at each phase:

python3 tools/infer.py -c configs/segmentation/unet_mobilenetv1_cityscapes.py --model-inputs img:${img-path} --save-path ${save_path}

Model Compilation

The quantization algorithms used in the HAT integrated quantitative training framework are specifically designed for Horizon processors, so fixed-point models trained with HAT can be compiled using the tools provided by HBDK to generate fixed-point models that can run on the computing platform. We provided the script compile_perf.py to facilitate the compilation process.

python3 tools/compile_perf.py -c configs/segmentation/unet_mobilenetv1_cityscapes.py --opt 2

Thus, we have got a segmentation task fixed-point model that can run on the Horizon processors from scratch.

Export FixedPoint Model

Once you've completed your quantization training, you can start exporting your fixed-point model. You can export it with the following command:

python3 tools/export_hbir.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Model Verification

After completing the training, we get the trained floating-point, quantitative, or fixed-point model. Similar to the training method, we can use the same method to complete metrics validation on the trained model and get the metrics of Float, Calibration, and Quantized, which are floating-point, quantitative, and fully fixed-point metrics, respectively.

python3 tools/predict.py --stage "float" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

python3 tools/predict.py --stage "calibration" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Similar to the model training, we can use --stage followed by "float", "calibration", to validate the trained floating-point model, and quantitative model, respectively.

The following command can be used to verify the accuracy of a fixed-point model, but it should be noted that hbir must be exported first:

python3 tools/predict.py --stage "int_infer" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Model Inference

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

python3 tools/infer_hbir.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py --model-inputs img:${img-path} --save-path ${save_path}

Simulation Board Accuracy Verification

In addition to the above model validation, we provide an accuracy validation method identical to the on-board environment, which can be accomplished by:

python3 tools/validation_hbir.py --stage "align_bpu" --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Fixed-point Model Checking and Compilation

As the quantitative training toolchain integrated in HAT is mainly prepared for Horizon's processors, it is a must to check and compile the quantitative models.

We provide an interface for model checking in HAT, which allows the user to define a quantitative model and then check whether it can work properly on the BPU first.

python3 tools/model_checker.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

After the model is trained, you can use the compile_perf_hbir script to compile the quantitative model into an HBM file that supports on-board running. The tool can also predict the performance on the BPU.

python3 tools/compile_perf_hbir.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

The above is the whole process from data preparation to the generation of quantitative and deployable models.

Training Details

Model Structure

model = dict(
    type="Segmentor",
    backbone=dict(
        type="MobileNetV1",
        num_classes=-1,
        bn_kwargs=dict(eps=2e-5, momentum=0.1),
        alpha=0.25,
        dw_with_relu=True,
        include_top=False,
        flat_output=False,
    ),
    neck=dict(
        type="DwUnet",
        base_channels=8,
        bn_kwargs=dict(eps=2e-5, momentum=0.1),
        act_type=nn.ReLU,
        use_deconv=False,
        dw_with_relu=True,
        output_scales=(4, 8, 16, 32, 64),
    ),
    head=dict(
        type="SegHead",
        num_classes=19,
        in_strides=(4, 8, 16, 32, 64),
        out_strides=(4, 8, 16, 32, 64),
        stride2channels={
            4: 32,
            8: 64,
            16: 128,
            32: 256,
            64: 512
        },
        feat_channels=(32, 64, 128, 256, 512),
        stacked_convs=0,
        argmax_output=False,
        dequant_output=True,
        int8_output=True,
        upscale=False,
        output_with_bn=True,
        bn_kwargs=dict(eps=2e-5, momentum=0.1),
    ),
    losses=dict(
        type="ClassificationFocalLoss",
        loss_name="Focal",
        num_classes=19,
        weight=tuple(np.array((256, 128, 64, 32, 16)) / 19),
        reduction="mean",
    ),
)

The segmentation model mainly consists of backbone, neck, and head. Here we use MobileNetV1_0.25 as the backbone, which is a lightweight and efficient network structure. The neck uses the Unet structure, which can combine the featuremap at each scale and preserve the fine spatial information. The head is a convolutional layer responsible for the output of the final segmentation results.

We use FocalLoss as the loss function. FocalLoss can be regarded as a cross-entropy loss function with dynamic weights, which can better solve the problem of training difficulties caused by category imbalance.

The hierarchical structure of Unet is consistent with the idea of FPN and can be treated with the same training method, that is, constructing an output on each scale of Unet, and construct a loss function with this output and the ground truth with corresponding size for training, which supervises each network scale, provides richer reference information for network training, reduces the training difficulty, and improves the training speed and the final accuracy.

Meanwhile, considering that the final result we need is the max-sized network output (scale=4), to avoid the over large gradient of other sizes affecting the accuracy of the max-sized output, we add weights to the loss function according to scale. A layer with larger scale has smaller weight.

Usually, after defining a model, especially for some public models, we would have the need to check the FLOPs. HAT provides calops.py to calculate the operations of the model, which is implemented as below:

python3 tools/calops.py --config configs/segmentation/unet_mobilenetv1_cityscapes.py

Such ops-counting tool can support both floating-point and fixed-point models at the same time.

Data Pre-processing

First, we use LabelRemap to re-map the data labels to the interval [0, 18].

For the training set, SegRandomAffine can perform a random affine transformation on the image for data enhancement, and we only configure a random scaling without any rotation operation.

Since the training uses an FPN-like approach, we need to scale the labels to different sizes to train the models with different scales.

The final model running on the BPU uses the YUV444 format as the image input, while regular trainings use the RGB format. Therefore, HAT provides ImageBgrToYuv444 data enhancement to convert RGB data to the YUV444 format.

Finally, normalization is necessary for deep learning model training.

Note that here we use MultiBatchProcessor for the task training. This Processor supports data pre-processing on the GPU in batches. Since the data preprocessing of the segmentation task is relatively complex, bottlenecks will occur if we use CPU for the processing. However, using GPU will increase the video memory usage, resulting in a lower maximum batch size. That said, the final training speed is considerably improved.

data_loader = dict(
    type=torch.utils.data.DataLoader,
    dataset=dict(
        type="CityscapesFromLMDB",
        data_path=os.path.join(data_rootdir, "train_lmdb"),
        transforms=[
            dict(type="PILToTensor"),
        ],
        num_samples=2975,
    ),
    sampler=dict(type=torch.utils.data.DistributedSampler),
    batch_size=batch_size_per_gpu,
    shuffle=True,
    num_workers=dataloader_workers,
    pin_memory=True,
)
batch_processor = dict(
    type="MultiBatchProcessor",
    need_grad_update=True,
    batch_transforms=[
        dict(type="LabelRemap", mapping=CITYSCAPES_LABLE_MAPPINGS),
        dict(type="SegOneHot", num_classes=num_classes),
        dict(type="SegResize", size=data_shape[1:]),
        dict(
            type="SegRandomAffine",
            degrees=0,
            scale=(0.5, 2.0),
            interpolation=InterpolationMode.BILINEAR,
            label_fill_value=0,
        ),
        dict(type="ImageBgrToYuv444", rgb_input=True),
        dict(type="ImageNormalize", mean=128.0, std=128.0),
        dict(
            type="Scale",
            scales=tuple(1 / np.array(train_scales)),
            mode="bilinear"
        ),
    ],
    loss_collector=loss_collector,
)

Compared with the training set, the data preprocessing of the validation set does not require random affine transformation and multi-scale scaling. The other stages are the same and will not be repeated.

Training Strategies

The segmentation task can be regarded as a pixel-level classification task, so the training strategy is highly similar to that of the classification task.

First, the training speed can be increased by increasing the learning rate as much as possible while ensuring convergence.

When training with a certain learning rate until the accuracy no longer grows, you can reduce the learning rate appropriately, the model will continue to converge, and the accuracy can be further improved.

After the final training, compare the accuracy of the test set with the training set, and if the accuracy of the training set is too much higher than the test set, the model is considered to be overfitted. In this case, increasing the weight decay can enhance the generalization ability of the model, reduce the overfitting, and obtain a higher accuracy of the test set.

float_trainer = dict(
    type="distributed_data_parallel_trainer",
    model=model,
    data_loader=data_loader,
    optimizer=dict(
        type=torch.optim.SGD,
        params={"weight": dict(weight_decay=5e-3)},
        lr=0.01,
        momentum=0.9,
    ),
    batch_processor=batch_processor,
    device=None,
    num_epochs=train_epochs,
    callbacks=[
        stat_callback,
        dict(
            type="StepDecayLrUpdater",
            lr_decay_id=[200, 240, 280],
            lr_decay_factor=0.1,
        ),
        metric_updater,
        tb_callback,
        aidi_tb_callback,
        tb_loss_callback,
        aidi_tb_loss_callback,
        val_callback,
        ckpt_callback,
    ],
    sync_bn=True,
)

Quantitative Training Strategies

The purpose of quantitative training is to simulate the process of fix-point computation by performing simulated data quantization based on the trained floating-point model to minimize the accuracy loss during the quantitative to fixed-point model conversion.

Since the floating-point model has converged to a good state after sufficient training, the model usually needs only a little fine-tuning in the quantitative training. The learning rate should not be set too large. You can start trying from 1e-4 order of magnitude, leaving other parameters the same as the floating-point training. Similarly, as the model fine-tuning requires less amount of training, usually a few dozen training epochs is enough.

Since the model already has a good accuracy before quantitative training, room for the accuracy improvement is small. The simulated quantization of the data will also lead to large fluctuations in the training process. At this point, we need to patiently and carefully observe the trend from the fluctuations and adjust the parameters appropriately to get the best results.

Model Checking & Compilation

For HAT, the significance of the quantitative model is that it can be run directly on the BPU. Therefore, it is necessary to check and compile the quantitative model.

The compile_perf.py script provided by HAT first checks the model to ensure that it can run properly on the BPU. Then it compiles the model to make it runnable on the BPU. After that, the script will also test the compiled fixed-point model to predict its performance on the BPU. Run the command as below:

python3 tools/compile_perf.py -c configs/segmentation/unet_mobilenetv1_cityscapes.py --opt 2

Pre-trained Model

HAT has provided pre-trained models for this example, which are all included in the release package.

The hb_compile Tool

QAT

Model Export

Horizon Operator

Vision Algorithm

Sample

VP API Introduction

Data Structure

Function Interface

Model Inference API Introduction

Data Type and Data Structure

API Interface

Model Inference Tool Introduction

High Performance Algorithm

HPL API Introduction

Data Structure

Functional Interface

Data Structure

Functional Interface

FRAMEWORK

TUTORIALS

EXAMPLES

API REFERENCE

#UNet Segmentation Model Training

#Training Process

#Downloading Dataset

#Packing Dataset

#Model Training

#Metrics Validation

#Simulation of On-board Accuracy Validation

#Model Inference

#Model Compilation

#Export FixedPoint Model

#Model Verification

#Model Inference

#Simulation Board Accuracy Verification

#Fixed-point Model Checking and Compilation

#Training Details

#Model Structure

#Data Pre-processing

#Training Strategies

#Quantitative Training Strategies

#Model Checking & Compilation

#Pre-trained Model

UNet Segmentation Model Training

Training Process

Downloading Dataset

Packing Dataset

Model Training

Metrics Validation

Simulation of On-board Accuracy Validation

Model Inference

Model Compilation

Export FixedPoint Model

Model Verification

Model Inference

Simulation Board Accuracy Verification

Fixed-point Model Checking and Compilation

Training Details

Model Structure

Data Pre-processing

Training Strategies

Quantitative Training Strategies

Model Checking & Compilation

Pre-trained Model