通过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/tensorflow2ncnngithub.com

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

onnx/tensorflow-onnxgithub.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.pygithub.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/ncnngithub.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 模型

Netronlutzroeder.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专栏。

是tf2 不是tf

mlir 的意义

导出 mlir

编译 mlir2ncnn

转换 ncnn 模型

推荐阅读

目录