Model Quantization

Introduction

TPU-MLIR is the compiler project for the Sophgo deep-learning processor. It provides a complete toolchain that converts pretrained neural networks from different frameworks into bmodel files that run efficiently on the Sophgo intelligent-vision deep-learning processor. The source code is open on GitHub: https://github.com/sophgo/tpu-mlir .

The paper https://arxiv.org/abs/2210.15016 describes the overall design of TPU-MLIR.

The overall TPU-MLIR architecture is shown below:

_images/framework.png

The currently directly supported frameworks are ONNX, Pytorch, Caffe, and TFLite. Models from other frameworks must first be converted to ONNX. For how to convert models from other deep-learning frameworks to ONNX, see the ONNX official site: https://github.com/onnx/tutorials.

Model conversion is performed inside a specific Docker container, mainly in two steps: first model_transform.py converts the original model to an MLIR file, then model_deploy.py converts the MLIR file to a bmodel.

To convert to an INT8 model, call run_calibration.py to generate a calibration table and pass it to model_deploy.py.

If the INT8 model does not meet accuracy requirements, you can call run_qtable.py to generate a quantization table that decides which layers use floating-point computation, and then pass it to model_deploy.py to produce a mixed-precision model.

1. Set Up the TPU-MLIR Environment

1.1 Base Environment

To relieve storage pressure on the board, use a non-BM1684X Linux system (here using WSL as an example) for model quantization and conversion. If your environment satisfies python >= 3.10 and ubuntu:22.04, you can skip the Docker environment configuration (this subsection).

Because the libc version affects model conversion and quantization, we use the official image to set up the environment. TPU-MLIR is developed inside Docker; once Docker is configured, it can be compiled and run.

If you are using Docker for the first time, run the following commands to install and configure it (only needed once):

sudo apt install docker.io
sudo systemctl start docker
sudo systemctl enable docker
sudo groupadd docker
sudo usermod -aG docker $USER
newgrp docker

Pull the required image from Docker Hub:

docker pull sophgo/tpuc_dev:latest

If the pull fails, you can download the image directly with wget:

#使用wget下载所需的镜像
wget https://sophon-assets.sophon.cn/sophon-prod-s3/drive/25/04/15/16/tpuc_dev_v3.4.tar.gz
#加载镜像
docker load -i tpuc_dev_v3.4.tar.gz

Start the image environment:

#首次创建tpumlir环境使用下面命令，--name tpumlir这里名字可自定义设置
docker run --privileged --name tpumlir -v $PWD:/workspace -it sophgo/tpuc_dev:latest
#非首次创建直接使用下面命令
docker run -v $PWD:/workspace -it sophgo/tpuc_dev:latest

1.2 Install TPU-MLIR

TPU-MLIR provides the following three installation methods.

(1) Download and install directly from PyPI (recommended):

pip install tpu_mlir -i https://pypi.tuna.tsinghua.edu.cn/simple

(2) Download the latest tpu_mlir-*-py3-none-any.whl from TPU-MLIR GitHub, then install with pip:

pip install tpu_mlir-*-py3-none-any.whl

Tips

TPU-MLIR requires different dependencies when processing models from different frameworks. For models generated by ONNX or Torch, install the additional dependency environment with:

pip install tpu_mlir[onnx] -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install tpu_mlir[torch] -i https://pypi.tuna.tsinghua.edu.cn/simple

Five configurations are currently supported: onnx, torch, tensorflow, caffe, paddle. You can install multiple configurations with a single command, or install all dependencies:

pip install tpu_mlir[onnx,torch,caffe] -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install tpu_mlir[all] -i https://pypi.tuna.tsinghua.edu.cn/simple

(3) If you obtained a release package of the form tpu-mlir_${version}-${hash}-${date}.tar.gz (this package can be obtained by downloading sophon-SDK and looking in its subdirectory, typically under SDK-23.09-LTS-SP4\tpu-mlir_20231116_054500), configure it as follows:

#可通过下面命令选择下载SDK
wget https://sophon-assets.sophon.cn/sophon-prod-s3/drive/24/12/31/10/SDK-23.09-LTS-SP4.zip
# 如果此前有通过pip安装过mlir，需要卸载掉
pip uninstall tpu_mlir
#解压安装发布包
tar xvf tpu-mlir_${version}-${hash}-${date}.tar.gz
cd tpu-mlir_${version}-${hash}-${date}
source envsetup.sh #配置环境变量

It is recommended to use the TPU-MLIR image only for compiling and quantizing models; program compilation and execution should be done in the development and runtime environments. For more TPU-MLIR tutorials, see the related page.

2. Compile the Model

This section uses yolov5s.onnx as an example to show how to compile and port an ONNX model to run on the BM1684X platform. Other models can refer to the related examples.

2.1 Configure the Project Directory

Download tpu-mlir-resource.tar from the Assets on GitHub and decompress it; after decompression, rename the folder to tpu_mlir_resource:

#可手动下载，也可通过下面命令使用wget下载，推荐手动
wget https://github.com/sophgo/tpu-mlir/releases/download/v1.20/tpu-mlir-resource.tar
#解压工程目录
tar -xvf tpu-mlir-resource.tar
#修改文件名称
mv regression/ tpu-mlir-resource/

Tips

tpu-mlir-resource.tar is a sample resource file. If you want to convert your own model, this file is not required; see the developer manual for the relevant configuration.

Create a model_yolov5s_onnx directory and put both the model file and image files into model_yolov5s_onnx.

Steps:

mkdir model_yolov5s_onnx && cd model_yolov5s_onnx
wget https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.onnx
cp -rf tpu_mlir_resource/dataset/COCO2017 .
cp -rf tpu_mlir_resource/image .
mkdir workspace && cd workspace

2.2 ONNX to MLIR

If the model takes an image as input, before converting the model we need to understand its preprocessing. If the model uses a preprocessed npz file as input, preprocessing is not a concern.

The preprocessing can be expressed by the following formula (where x is the input):

$$ y=(x-mean)*scale $$

The official YOLOv5 image is in RGB format; each value is multiplied by 1/255, corresponding to mean and scale of 0.0,0.0,0.0 and 0.0039216,0.0039216,0.0039216.

The model conversion command is as follows:

$ model_transform \
    --model_name yolov5s \
    --model_def ../yolov5s.onnx \
    --input_shapes [[1,3,640,640]] \
    --mean 0.0,0.0,0.0 \
    --scale 0.0039216,0.0039216,0.0039216 \
    --keep_aspect_ratio \
    --pixel_format rgb \
    --output_names 350,498,646 \
    --test_input ../image/dog.jpg \
    --test_result yolov5s_top_outputs.npz \
    --mlir yolov5s.mlir

The main parameters of model_transform are described below (for the full description, see the TPU-MLIR developer reference — User Interface chapter):

Parameter	Required?	Description
model_name	Yes	Model name
model_def	Yes	Model definition file, e.g. `.onnx`, `.tflite`, or `.prototxt`
input_shapes	No	Input shape, e.g. `[[1,3,640,640]]`; a 2D array that supports multiple inputs
input_types	No	Input type, e.g. int32; separate multiple inputs with `,`; defaults to float32
resize_dims	No	The size to resize the original image to; if not specified, resized to the model input size
keep_aspect_ratio	No	Whether to keep aspect ratio on resize; default false; when set, pads the shortfall with 0
mean	No	Per-channel mean of the image; default `0.0,0.0,0.0`
scale	No	Per-channel scale of the image; default `1.0,1.0,1.0`
pixel_format	No	Pixel format: rgb, bgr, gray, or rgbd; default bgr
channel_format	No	Channel format: nhwc or nchw for image input, none for non-image input; default nchw
output_names	No	Output names; if not specified, the model's outputs are used; when specified, these names are used as outputs
test_input	No	Input file for validation; can be an image, npy, or npz; if not specified, no correctness validation is performed
test_result	No	Output file after validation
excepts	No	Names of network layers to exclude from validation; multiple separated by `,`
mlir	Yes	Output MLIR file name and path

After conversion to MLIR, a ${model_name}_in_f32.npz file is generated; this is the model's input file.

2.3 MLIR to F16 Model

Convert the MLIR file to an F16 bmodel as follows:

model_deploy \
    --mlir yolov5s.mlir \
    --quantize F16 \
    --processor bm1684x \
    --test_input yolov5s_in_f32.npz \
    --test_reference yolov5s_top_outputs.npz \
    --model yolov5s_1684x_f16.bmodel

After compilation, a file named yolov5s_1684x_f16.bmodel is generated.

The main parameters of model_deploy are described below (for the full description, see the TPU-MLIR developer reference — User Interface chapter):

Parameter	Required?	Description
mlir	Yes	The MLIR file
quantize	Yes	Default quantization type; supports F32/F16/BF16/INT8
processor	Yes	Target platform; supports bm1690, bm1688, bm1684x, bm1684, cv186x, cv183x, cv182x, cv181x, cv180x
calibration_table	No	Calibration table path; required when INT8 quantization is present
tolerance	No	Error tolerance between the MLIR quantized result and the MLIR fp32 inference result
test_input	No	Input file for validation; can be an image, npy, or npz; if not specified, no correctness validation is performed
test_reference	No	Reference data for validating model correctness (npz format); the per-op computation results
compare_all	No	Whether to compare all intermediate results during validation; by default intermediate results are not compared
excepts	No	Names of network layers to exclude from validation; multiple separated by `,`
op_divide	No	cv183x/cv182x/cv181x/cv180x only; tries to split large ops into smaller ones to save ion memory; applies to a few specific models
model	Yes	Output model file name and path
num_core	No	When target is bm1688, selects the number of TPU cores for parallel computation; default 1 TPU core
skip_validation	No	Skip bmodel correctness validation to speed up deployment; validation runs by default

2.5 MLIR to INT8 Model

2.5.1 Generate the Calibration Table

Before converting to an INT8 model, run calibration to obtain a calibration table; prepare about 100–1000 input images as appropriate.

Then use the calibration table to generate a symmetric or asymmetric bmodel. If symmetric meets the requirement, asymmetric is generally not recommended because its performance is slightly worse than the symmetric model.

Here we use 100 existing images from COCO2017 as an example to run calibration:

run_calibration yolov5s.mlir \
    --dataset ../COCO2017 \
    --input_num 100 \
    -o yolov5s_cali_table

After running, a file named yolov5s_cali_table is generated; this file is the input for subsequent INT8 model compilation.

2.5.2. Compile to an INT8 Symmetric Quantization Model

To convert to an INT8 symmetric quantization model, run:

model_deploy \
    --mlir yolov5s.mlir \
    --quantize INT8 \
    --calibration_table yolov5s_cali_table \
    --processor bm1684x \
    --test_input yolov5s_in_f32.npz \
    --test_reference yolov5s_top_outputs.npz \
    --tolerance 0.85,0.45 \
    --model yolov5s_1684x_int8_sym.bmodel

After compilation, a file named yolov5s_1684x_int8_sym.bmodel is generated.

2.6 Result Comparison

This release includes a YOLOv5 use case written in Python; use the detect_yolov5 command to perform object detection on an image.

The source path for this command is {package/path/to/tpu_mlir}/python/samples/detect_yolov5.py.

Reading that code shows how the model is used: first preprocess to get the model input, then inference to get the output, and finally post-process.

The following commands verify the ONNX/F16/INT8 results respectively.

Run the ONNX model as follows to get dog_onnx.jpg:

detect_yolov5 \
    --input ../image/dog.jpg \
    --model ../yolov5s.onnx \
    --output dog_onnx.jpg

onnx

Run the F16 bmodel as follows to get dog_f16.jpg:

detect_yolov5 \
    --input ../image/dog.jpg \
    --model yolov5s_1684x_f16.bmodel \
    --output dog_f16.jpg

f16 bmodel

Run the INT8 symmetric bmodel as follows to get dog_int8_sym.jpg:

detect_yolov5 \
    --input ../image/dog.jpg \
    --model yolov5s_1684x_int8_sym.bmodel \
    --output dog_int8_sym.jpg

onnx