概述

NNTC支持TFLite量化模型直接转换,这里以ResNet50模型举例,说明编译方法

本例模型ResNet50.tflite下载地址: https://disk.sophgo.vip/sharing/ibo9YyxgT

TFLite模型编译用到了TPU-NNTC里的bmtflite前端

初始化tpu-nntc环境

在 docker 环境使用工具链软件,最新版本的 docker 可以参考 官方教程 进行安装。安装完成后,执行下面的脚本将当前用户加入 docker 组,获得 docker 执行权限。

sudo usermod -aG docker $USER
newgrp docker

从官网上下载tpu-nntc的压缩包,命名如 tpu-nntc_vx.y.z-<hash>-<date>.tar.gz

mkdir tpu-nntc
# 将压缩包解压到tpu-nntc
tar zxvf tpu-nntc_vx.y.z-<hash>-<date>.tar.gz --strip-components=1 -C tpu-nntc

tpu-nntc使用的docker是sophgo/tpuc_dev:v2.1, docker镜像和tpu-nntc有绑定关系,少数情况下有可能更新了tpu-nntc,需要新的镜像

cd tpu-nntc

# 进入docker,如果当前系统没有对应镜像,会自动从docker hub上下载
# 将tpu-nntc的上一级目录映射到docker内的/workspace目录
# 这里用了8001到8001端口映射,之后在使用ufw可视化工具会用到
# 如果端口已经占用,请更换其他未占用端口,后面根据需要更换进行调整
docker run -v $PWD/..:/workspace -p 8001:8001 -it sophgo/tpuc_dev:v2.1

# 此时已经进入docker,并在/workspace目录下
# 下面初始化软件环境
cd /workspace/tpu-nntc
source scripts/envsetup.sh

其他docker初始化方式:

如果是离线环境,可以先提前下载,然后导入到目标机上,命令如下:

# 在可以联网的机器上, 保存镜像
docker pull sophgo/tpuc_dev:v2.1
docker save sophgo/tpuc_dev:v2.1 -o sophon_tpuc_dev.docker

# 将sophon_tpuc_dev_docker.tar复制到目标机器上后
docker load -i sophon_tpuc_dev.docker

此外,tpu-nntc在scripts目录提供了docker_setup.sh脚本, 供参考和方便基本使用,实际使用中可根据需求自行管理。使用方法如下:

# docker_setup.sh 后面参数来指定哪个目录作为工作目录,映射到docker里的/workspace
# 这里用 tpu-nntc 的上一级目录
scripts/docker_setup.sh ..

# 会输出以下信息,如果8001端口占用,会自动选择其他未占用端口
# ...
# OPEN 'http://localhost:8001' IN WEB BROWSER WHEN 'python3 -m ufw.tools.app' RUNNING IN DOCKER

编译BMODEL

获取模型参数

可以netron或其他软件打开ResNet50.tflite文件,得到input_names和shapes参数

../_images/resnet50_in.png

在命令行下编译

# ResNet50.tflite 需要在当前目录下
# 由于转换的是TFLite模型,所以需要bmtflite前端
# 这里只有一个输入,如果有多个,可以用","隔开
python3 -m bmtflite --model ResNet50.tflite \
                --input_names "serving_default_input_1:0" \
                --shapes "[1,224,224,3]" \
                --target BM1684X \
                --outdir cmd_bmodel

查看下bmodel的信息:

tpu_model --info cmd_bmodel/compilation.bmodel

输出模型信息如下:

bmodel version: B.2.2
chip: BM1684X
create time: Mon Aug 22 21:12:23 2022

==========================================
net 0: [network]  static
------------
stage 0:
input: serving_default_input_1:0, [1, 224, 224, 3], float32, scale: 1
output: StatefulPartitionedCall:0, [1, 1000], float32, scale: 1

device mem size: 28124160 (coeff: 26162176, instruct: 0, runtime: 1961984)
host mem size: 0 (coeff: 0, runtime: 0)

用Python脚本编译

创建python编译脚本compile_resnet50.py, 内容如下

import bmtflite as bm

bm.compile(
  model = "ResNet50.tflite",
  input_names = ["serving_default_input_1:0"],
  shapes = [[1,224,224,3]],
  target = "BM1684X",
  outdir = "py_bmodel",
  net_name = "resnet50",
  )

执行编译

# ResNet50.tflite 需要在当前目录下
python3 compile_resnet50.py

查看下bmodel的信息:

tpu_model --info py_bmodel/compilation.bmodel

输出如下:

bmodel version: B.2.2
chip: BM1684X
create time: Mon Aug 22 21:19:34 2022

==========================================
net 0: [resnet50]  static
------------
stage 0:
input: serving_default_input_1:0, [1, 224, 224, 3], float32, scale: 1
output: StatefulPartitionedCall:0, [1, 1000], float32, scale: 1

device mem size: 28124160 (coeff: 26162176, instruct: 0, runtime: 1961984)
host mem size: 0 (coeff: 0, runtime: 0)