通过MLIR将tensorflow2模型转换到ncnn

作者:nihui
转自:知乎

是tf2 不是tf

请注意!请注意!这里说的是 tensorflow2,不是 tensorflow,这两个不是一个东西,以下内容只适用于 tensorflow2。

曾经,我写过 tensorflow2ncnn.cpp,从 tensorflow frozen pb 转成 ncnn param+bin。

然而,我真的没用过 tensorflow,这转换器就是半猜半试写出来的,各种 bug,最后只好承认自己菜,放弃了,删掉了。 tensorflow 那么乱七八糟的 API,我不想学,一点都不想看。

后来,有社区大佬 hanzy88 接手这份代码,并把 vgg16 resnet yolov3 弄好了,ncnn 置顶多年的 issue 终于有了着落。

https://github.com/hanzy88/tensorflow2ncnn​github.com

好多人以为 ncnn 支持 onnx 模型,就用 tensorflow-onnx 转换器,先转成 onnx,结果跑 onnx2ncnn 各种报错

onnx/tensorflow-onnx​github.com

ncnn 所支持的 onnx,是 pytorch 转来的onnx,不是别的框架转来的 onnx。不一样!不一样!

mlir 的意义

tensorflow2 出来,据说是完全全新的 API,我打开 tf2 官网看着 tutorial 学习了下。 我之前一直用 mxnet/gluon 捏模型,tf2 作为后来者,体验了下 API,感觉 tf2 完成了 gluon 所没有完成的所有技术梦想,好感度+++。

tf2 里头,frozen pb 是 deprecated 的东西了,mlir 是交换模型的正统选择

我搞不懂为何很多人觉得 mlir 和 tvm 有联系,在我眼里,mlir 就是 google 版的 onnx,就是一种后缀为 .mlir 的模型格式,tf2 模型可以导出成 mlir

合理推测,未来 tensorflow-onnx 项目会长期处于社区维护状态,而不会合入 tensorflow 内置。mlir 作为 onnx 的直接竞品和可能的替代者,google 必然强力推广,llvm-11 开始自带 mlir,分发到各个 Linux 系统里。

本着方便自己也方便大家的原则,于是我就写了个 mlir2ncnn 转换器,支持 tf2 转来的 mlir,不是别的框架转来的 mlir。不一样!不一样!

下面详细介绍下 tf2-mlir-ncnn 整个流程,这里我使用 tf2 官网上的 pix2pix 例子

https://github.com/tensorflow/examples/blob/master/tensorflow\_examples/models/pix2pix/pix2pix.py​github.com

导出 mlir

给 pix2pix 增加个 save\_mlir 方法,导出 mlir

print ('Training ...')
  #return pix2pix_object.train(train_dataset, checkpoint_pr)
  pix2pix_object.save_mlir()

加载训练时保存的 checkpoint

def save_mlir(self):
    self.checkpoint.restore("./training_checkpoints/ckpt-100")

图优化,用的是 grappler 的功能,这块东西按照规划要放在 mlir pass 实现,但是 mlir 导出那边似乎优化做得并没有 grappler 好,那就继续 grappler 好了

self.generator 就是 pix2pix 的主模型,get\_concrete\_function 里面的是输入 shape,这个 shape 一定要写清楚,不然会因为不知道输入 shape,转出来的 mlir 采用动态 shape,搞出来一大堆胶水 op

get\_grappler\_config 后面的列表是我在 tf2 代码里翻出来的一些实用的 pass 名称。还有个叫 “layout” 的 pass,会在 convolution 前后增加 transpose 胶水 op,不用它

from tensorflow.lite.python.util import run_graph_optimizations, get_grappler_config
    from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph

    fff = tf.function(self.generator).get_concrete_function(tf.TensorSpec([1, 256, 256, 3], tf.float32))

    frozen_func, graph_def = convert_variables_to_constants_v2_as_graph(fff)

    input_tensors = [
        tensor for tensor in frozen_func.inputs
        if tensor.dtype != tf.resource
    ]
    output_tensors = frozen_func.outputs

    graph_def = run_graph_optimizations(
        graph_def,
        input_tensors,
        output_tensors,
        config=get_grappler_config(['pruning', 'function', 'constfold', 'shape', 'remap', 'memory', 'common_subgraph_elimination', 'arithmetic', 'loop', 'dependency', 'debug_stripper']),
        graph=frozen_func.graph)

导出mlir,tf.mlir.experimental.convert\_graph\_def 有个 pass 参数可以做一些图优化的事情,前面已经优化好了就不麻烦了

tf_mlir_graph = tf.mlir.experimental.convert_graph_def(graph_def)

    outfile = open('pix2pix.mlir', 'wb')
    outfile.write(tf_mlir_graph.encode())
    outfile.close()

编译 mlir2ncnn

mlir 本体在 llvm 大城堡里,但是 llvm-10.0 没有,一定要用最新的 git 版本才能编译出 mlir

git clone --depth=1 https://github.com/llvm/llvm-project.git

很久很久以后,开始编译 mlir,注意用 Release 编译,打开动态库,不然编译出来二进制会有十几个G那么大,不仅没用,还非常浪费空间

cd llvm-project
mkdir build
cd build
cmake -DCMAKE_INSTALL_PREFIX=install -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=ON -DLLVM_ENABLE_PROJECTS="mlir" -DLLVM_INCLUDE_EXAMPLES=OFF -DLLVM_INCLUDE_TESTS=OFF ../llvm/
make -j8
make install

mlir2ncnn 基本功能实现的差不多了,今天在 pr 里面,也许过一些天就在 master 上了。 转换器代码中的 td 和一些 cpp 是从 tensorflow 代码里挖出来的,用来解析 tf2 dialect,这也是为什么这个转换器只支持 tf2 导出的mlir。

[WIP] MLIR converter by nihui · Pull Request #1810 · Tencent/ncnn​github.com

编译前,需要改下 CMakeLists.txt 里的 llvm 和 mlir 安装目录

set(LLVM_DIR "/home/nihui/osd/llvm-project/build/install/lib/cmake/llvm")
find_package(LLVM REQUIRED)

set(MLIR_DIR "/home/nihui/osd/llvm-project/build/install/lib/cmake/mlir")
find_package(MLIR REQUIRED)

常规编译就好

cd tools/mlir
mkdir build
cd build
cmake ..
make

转换 ncnn 模型

netron visualizer 截至发稿时,还不支持可视化 mlir 模型

Netron​lutzroeder.github.io

pix2pix是个普通的unet模型,可以无痛转换

./mlir2ncnn pix2pix.mlir pix2pix.param pix2pix.bin

针对 mlir 转出来模型的特点,如 convolution-bias 分开等,我额外补充了一些图优化功能,用 ncnnoptimize 可以优化掉

./ncnnoptimize pix2pix.param pix2pix.bin pix2pix-opt.param pix2pix-opt.bin 65536

啊!强迫症非常满足!

mlir 模型有个特点,就是 op 和 blob 名字是不存在的,只有上下文关系,为了区分,我用了 “文件名:行:列” 组成唯一标识符,有点像编译错误的那种输出(

ncnn::Mat in;
ex.input("pix2pix.mlir:52:11", in);

ncnn::Mat out;
ex.extract("pix2pix.mlir:114:12", out);

更多嵌入式AI算法部署等请关注极术嵌入式AI专栏

发表评论

邮箱地址不会被公开。 必填项已用*标注

Related Post