Awesome

TPU-MLIR

For Chinese version: README.

TPU-MLIR is an open-source machine-learning compiler based on MLIR for TPU. This project provides a complete toolchain, which can convert pre-trained neural networks from different frameworks into binary files bmodel that can be efficiently operated on TPUs.

SOPHGO aims to become a leading global provider of general-purpose computing power. SOPHGO focuses on the research, development, and promotion of computing products such as Deep Learning and RISC-V processors, and has built a comprehensive application matrix covering the 'cloud, edge, and endpoint' scenarios with its self-developed products. SOPHGO provides computing products and integrated solutions for applications such as smart cities, intelligent computing centers, smart security, intelligent transportation, safety production, industrial quality inspection, and intelligent terminals. The company has research and development centers in more than 10 cities in China including Beijing, Shanghai, Shenzhen, Qingdao, and Xiamen, as well as in the United States and Singapore.

Currently, supported Deep Learning frameworks are PyTorch, ONNX, TFLite and Caffe. Models from other frameworks need to be converted to ONNX models.

Prebuilt TPU-MLIR Python Package

We provide TPU-MLIR python package for quick installation by skipping the building step. Environment prerequisite: python >= 3.10 and ubuntu:22.04 (It is recommended to use our docker image directly).

To install:

pip install tpu_mlir

Resources

Here are some resources to help you better understand the project:

If you have any questions while doing the tasks above, you can ask or check the existing answers in our Q&A Platform.

How to Build

After cloning the code of this project, it needs to be compiled in docker.

• Download the required image from dockerhub.

docker pull sophgo/tpuc_dev:latest

• If the pulling failed, you can download it in the following way:

wget https://sophon-file.sophon.cn/sophon-prod-s3/drive/24/06/14/12/sophgo-tpuc_dev-v3.2_191a433358ad.tar.gz

docker load -i sophgo-tpuc_dev-v3.2_191a433358ad.tar.gz

• Create the container:

# myname1234 is just an example, you can set your own name
docker run --privileged --name myname1234 -v $PWD:/workspace -it sophgo/tpuc_dev:latest

After the container is created, the directory of the code in docker should be /workspace/tpu-mlir.

• Building

Run the following command in the project directory:

cd tpu-mlir
source ./envsetup.sh
./build.sh

Usage

Introduce the usage of TPU-MLIR by a simple example of compiling yolov5s.onnx and running it on the BM1684X TPU platform.

The model comes from the official website of yolov5: https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.onnx.

It has been placed in project path regression/model/yolov5s.onnx.

Preparation

Firstly, create a model_yolov5s directory at the same level directory with this project. Then put both model and image files into it.

The operation is as follows:

mkdir model_yolov5s && cd model_yolov5s
cp ${REGRESSION_PATH}/model/yolov5s.onnx .
cp -rf ${REGRESSION_PATH}/dataset/COCO2017 .
cp -rf ${REGRESSION_PATH}/image .
mkdir workspace && cd workspace

Model to MLIR

If the model takes images as input, we need to learn its preprocessing before transforming. No preprocessing needs to be considered if the input is npz file. The preprocessing process is formulated as follows:

$$ y = （x - mean） \times scale, $$

where x represents the input.

The input of the official yolov5 is RGB image. Each value will be multiplied by 1/255. Mean and scale are 0.0, 0.0, 0.0 and 0.0039216, 0.0039216, 0.0039216 respectively.

The model conversion command:

model_transform.py \
    --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

Main arguments of model_transform.py (for complete information please check the technical reference manual):

After converting to mlir file, a ${model_name}_in_f32.npz file containing preprocessed input will be generated.

MLIR to F16 bmodel

Convert the mlir file to the F16 bmodel by the following command:

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

Main arguments of model_deploy.py (for complete information please check the technical reference manual):

MLIR to INT8 bmodel

Before converting to the INT8 model, you need to run calibration to get the calibration table. The number of input data is about 100 to 1000 according to the situation.

Then use the calibration table to generate a symmetric int8 bmodel. It is generally not recommended to use the asymmetric one if the symmetric one already meets the requirements, because the performance of the asymmetric model will be slightly worse than the symmetric model.

Here is an example of the existing 100 images from COCO2017 to perform calibration:

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

Execute the following command to convert to the INT8 symmetric quantized model:

model_deploy.py \
  --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.bmodel

Results Comparison

This project has a yolov5 sample written in python (path: python/samples/detect_yolov5.py) for object detection. Read the code to learn how the model is used:

1. preprocess the input
2. model inference to get output
3. post-process the output

The following code is used to verify the output of onnx/f32/int8 model respectively:

• ONNX model:

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

• F16 bmodel:

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

• INT8 symmetric quantized bmodel:

detect_yolov5.py \
  --input ../image/dog.jpg \
  --model yolov5s_1684x_int8.bmodel \
  --output dog_int8.jpg

Outputs of different models are compared below:

Auxiliary Tools

Model Inference Tool model_runner.py

Supports bmodel/mlir/pytorch/onnx/tflite/caffe.

model_runner.py \
  --input resnet18_in_f32.npz \
  --model resnet18_1684x_f32.bmodel \
  --output resnet18_output.npz

Tool for bmodel

The bmodel file can be viewed and edited by model_tool:

  model_tool
    --info model_file : show brief model info
    --print model_file : show detailed model info
    --extract model_file : extract one multi-net bmodel to multi one-net bmodels
    --combine file1 .. fileN -o new_file: combine bmodels to one bmodel by filepath
    --combine_dir dir1 .. dirN -o new_dir: combine bmodels to one bmodel by directory path
    --dump model_file start_offset byte_size out_file: dump binary data to file from bmodel

For example, to get basic information of bmodel:

model_tool --info resnet18_1684x_f32.bmodel

Related Links

• Official Network
• Discourse
• Videos