4. 编译TORCH模型

本章以 yolov5s.pt 为例, 介绍如何编译迁移一个pytorch模型至BM1684X TPU平台运行。

本章需要如下文件(其中xxxx对应实际的版本信息):

tpu-mlir_xxxx.tar.gz (tpu-mlir的发布包)

4.1. 加载tpu-mlir

以下操作需要在Docker容器中。关于Docker的使用, 请参考 启动Docker Container 。

1$ tar zxf tpu-mlir_xxxx.tar.gz
2$ source tpu-mlir_xxxx/envsetup.sh

envsetup.sh 会添加以下环境变量:

表 4.1 环境变量

变量名

值

说明

TPUC_ROOT

tpu-mlir_xxx

解压后SDK包的位置

MODEL_ZOO_PATH

${TPUC_ROOT}/../model-zoo

model-zoo文件夹位置, 与SDK在同一级目录

REGRESSION_PATH

${TPUC_ROOT}/regression

regression文件夹的位置

envsetup.sh 对环境变量的修改内容为:

 1 export PATH=${TPUC_ROOT}/bin:$PATH
 2 export PATH=${TPUC_ROOT}/python/tools:$PATH
 3 export PATH=${TPUC_ROOT}/python/utils:$PATH
 4 export PATH=${TPUC_ROOT}/python/test:$PATH
 5 export PATH=${TPUC_ROOT}/python/samples:$PATH
 6 export PATH=${TPUC_ROOT}/customlayer/python:$PATH
 7 export LD_LIBRARY_PATH=$TPUC_ROOT/lib:$LD_LIBRARY_PATH
 8 export PYTHONPATH=${TPUC_ROOT}/python:$PYTHONPATH
 9 export PYTHONPATH=${TPUC_ROOT}/customlayer/python:$PYTHONPATH
10 export MODEL_ZOO_PATH=${TPUC_ROOT}/../model-zoo
11 export REGRESSION_PATH=${TPUC_ROOT}/regression

4.2. 准备工作目录

建立 model_yolov5s_pt 目录, 注意是与tpu-mlir同级目录; 并把模型文件和图片文件都 放入 model_yolov5s_pt 目录中。

操作如下:

1$ mkdir model_yolov5s_pt && cd model_yolov5s_pt
2$ wget https://github.com/sophgo/model-zoo/raw/main/vision/detection/yolov5/yolov5s-5.0.pt
3$ cp -rf $TPUC_ROOT/regression/dataset/COCO2017 .
4$ cp -rf $TPUC_ROOT/regression/image .
5$ mkdir workspace && cd workspace

这里的 $TPUC_ROOT 是环境变量, 对应tpu-mlir_xxxx目录。

4.3. TORCH转MLIR

本例中的模型是 RGB 输入, mean和scale分别为 0.0,0.0,0.0 和 0.0039216,0.0039216,0.0039216。

模型转换命令如下:

$ model_transform.py \
    --model_name yolov5s_pt \
    --model_def ../yolov5s-5.0.pt \
    --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 \
    --test_input ../image/dog.jpg \
    --test_result yolov5s_pt_top_outputs.npz \
    --mlir yolov5s_pt.mlir

转成mlir文件后, 会生成一个 ${model_name}_in_f32.npz 文件, 该文件是模型的输入文件。

4.4. MLIR转F16模型

将mlir文件转换成f16的bmodel, 操作方法如下:

$ model_deploy.py \
    --mlir yolov5s_pt.mlir \
    --quantize F16 \
    --chip bm1684x \
    --test_input yolov5s_pt_in_f32.npz \
    --test_reference yolov5s_pt_top_outputs.npz \
    --model yolov5s_pt_1684x_f16.bmodel

编译完成后, 会生成名为 yolov5s_pt_1684x_f16.bmodel 的文件。

4.5. MLIR转INT8模型

4.5.1. 生成校准表

转INT8模型前需要跑calibration, 得到校准表; 这里用现有的100张来自COCO2017的图片举例, 执行calibration:

$ run_calibration.py yolov5s_pt.mlir \
    --dataset ../COCO2017 \
    --input_num 100 \
    -o yolov5s_pt_cali_table

运行完成后会生成名为 yolov5s_pt_cali_table 的文件, 该文件用于后续编译INT8 模型的输入文件。

4.5.2. 编译为INT8对称量化模型

转成INT8对称量化模型, 执行如下命令:

$ model_deploy.py \
    --mlir yolov5s_pt.mlir \
    --quantize INT8 \
    --calibration_table yolov5s_pt_cali_table \
    --chip bm1684x \
    --test_input yolov5s_pt_in_f32.npz \
    --test_reference yolov5s_pt_top_outputs.npz \
    --tolerance 0.85,0.45 \
    --model yolov5s_pt_1684x_int8_sym.bmodel

编译完成后, 会生成名为 yolov5s_pt_1684x_int8_sym.bmodel 的文件。

4.6. 效果对比

利用 $TPUC_ROOT/python/samples/detect_yolov5.py 路径下的源码, 对图片进行目标检测。 用以下代码分别来验证pytorch/f16/int8的执行结果。

pytorch模型的执行方式如下, 得到 dog_torch.jpg :

$ detect_yolov5.py \
    --input ../image/dog.jpg \
    --model ../yolov5s.pt \
    --output dog_torch.jpg

f16 bmodel的执行方式如下, 得到 dog_f16.jpg :

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

int8对称bmodel的执行方式如下, 得到 dog_int8_sym.jpg :

$ detect_yolov5.py \
    --input ../image/dog.jpg \
    --model yolov5s_pt_1684x_int8_sym.bmodel \
    --output dog_int8_sym.jpg

对比结果如下:

_images/yolov5s_pt.png

图 4.1 TPU-MLIR对YOLOv5s编译效果对比

由于运行环境不同, 最终的效果和精度与 图 4.1 会有些差异。