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 会添加以下环境变量:
变量名 |
值 |
说明 |
|---|---|---|
TPUC_ROOT |
tpu-mlir_xxx |
解压后SDK包的位置 |
MODEL_ZOO_PATH |
${TPUC_ROOT}/../model-zoo |
model-zoo文件夹位置, 与SDK在同一级目录 |
envsetup.sh 对环境变量的修改内容为:
1export PATH=${TPUC_ROOT}/bin:$PATH
2export PATH=${TPUC_ROOT}/python/tools:$PATH
3export PATH=${TPUC_ROOT}/python/utils:$PATH
4export PATH=${TPUC_ROOT}/python/test:$PATH
5export PATH=${TPUC_ROOT}/python/samples:$PATH
6export LD_LIBRARY_PATH=$TPUC_ROOT/lib:$LD_LIBRARY_PATH
7export PYTHONPATH=${TPUC_ROOT}/python:$PYTHONPATH
8export MODEL_ZOO_PATH=${TPUC_ROOT}/../model-zoo
9export REGRESSION_PATH=${TPUC_ROOT}/regression
4.2. 准备工作目录
建立 model_yolov5s_pt 目录, 注意是与tpu-mlir同级目录; 并把模型文件和图片文件都
放入 model_yolov5s_pt 目录中。
操作如下:
1$ mkdir yolov5s_torch && cd yolov5s_torch
2$ cp $TPUC_ROOT/regression/model/yolov5s.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.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 \
--tolerance 0.99,0.99 \
--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
对比结果如下:
图 4.1 TPU-MLIR对YOLOv5s编译效果对比
由于运行环境不同, 最终的效果和精度与 图 4.1 会有些差异。