6. 编译TFLite模型

本章以 lite-model_mobilebert_int8_1.tflite 模型为例, 介绍如何编译迁移一个TFLite模型至BM1684X TPU平台运行。

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

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

6.1. 加载tpu-mlir

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

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

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

表 6.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

6.2. 准备工作目录

建立 mobilebert_tf 目录, 注意是与tpu-mlir同级目录; 并把测试图片文件放入 mobilebert_tf 目录中。

操作如下:

1$ mkdir mobilebert_tf && cd mobilebert_tf
2$ wget -O lite-model_mobilebert_int8_1.tflite https://storage.googleapis.com/tfhub-lite-models/iree/lite-model/mobilebert/int8/1.tflite
3$ cp ${REGRESSION_PATH}/npz_input/squad_data.npz .
4$ mkdir workspace && cd workspace

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

6.3. TFLite转MLIR

模型转换命令如下:

$ model_transform.py \
    --model_name mobilebert_tf \
    --mlir mobilebert_tf.mlir \
    --model_def ../lite-model_mobilebert_int8_1.tflite \
    --test_input ../squad_data.npz \
    --test_result mobilebert_tf_top_outputs.npz \
    --input_shapes [[1,384],[1,384],[1,384]] \
    --channel_format none

转成mlir文件后, 会生成一个 mobilebert_tf_in_f32.npz 文件, 该文件是模型的输入文件。

6.4. MLIR转模型

该模型是tflite非对称量化模型, 可以按如下参数转成模型:

$ model_deploy.py \
    --mlir mobilebert_tf.mlir \
    --quantize INT8 \
    --asymmetric \
    --chip bm1684x \
    --test_input mobilebert_tf_in_f32.npz \
    --test_reference mobilebert_tf_top_outputs.npz \
    --model mobilebert_tf_bm1684x_int8_asym.bmodel

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