Tensorflow PB模型转换为OM模型

前言：Tensorflow训练出的PB模型，若想在Ascend 310芯片上运行，需转换成OM模型，其主要转换流程包括：1. Tensorflow SaveModel转换为Tensorflow FrozenGraph；2. 利用OMG或ATC工具，将Tensorflow FrozenGraph模型转换为OM模型。

1. Tensorflow SaveModel转换为FrozenGraph

Tensorflow框架可导出多种不同格式的模型，如CheckPoint、SaveModel、Frozen GraphDef、HDF5等，详情可参考Tensorflow模型导出总结 。通常而言，TF默认导出格式为SaveModel，参数以及网络结构是分开保存。

├── assets/ # 所需的外部文件，例如说初始化的词汇表文件，一般无
├── saved_model.pb # 保存的是MetaGraph的网络结构
├── variables # 参数权重，包含了所有模型的变量(tf.Variable objects)参数
    ├── variables.data-00000-of-00001
    └── variables.index

而转换为OM模型，需使用Ascend提供的转换工具，其输入格式为FrozenGraph，参数和模型结构是保存在同一个文件。

├── saved_model.pb # 保存的是MetaGraph的网络结构

因此，需将SaveModel转换为FrozenGraph，转换方法可通过Tensorflow提供的API，转换方法如下所示。

from tensorflow.python.tools import freeze_graph
from tensorflow.python.saved_model import tag_constants
input_saved_model_dir = "~/save_model/"
output_node_names = "loss/BiasAdd"
input_binary = False
input_saver_def_path = False
restore_op_name = None
filename_tensor_name = None
clear_devices = False
input_meta_graph = False
checkpoint_path = None
input_graph_filename = None
saved_model_tags = tag_constants.SERVING
output_graph_filename='~/frozen_graph.pb'
freeze_graph.freeze_graph(input_graph_filename,
  input_saver_def_path,
  input_binary,
  checkpoint_path,
  output_node_names,
  restore_op_name,
  filename_tensor_name,
  output_graph_filename,
  clear_devices,
  "", "", "",
  input_meta_graph,
  input_saved_model_dir,
  saved_model_tags)

其中，input_save_model_dir 为训练完成后导出的SaveModel路径，output_graph_filename是需要导出的frozen_graph模型文件名，output_node_names是模型输出的最后一层名称，可通过saved_model_cli命令查看。

$ saved_model_cli show --all --dir ~/save_model_dir/

# 输出信息
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['predict_object']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['input_ids'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 500)
        name: input_ids:0
    inputs['input_mask'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 500)
        name: input_mask:0
    inputs['segment_ids'] tensor_info:
        dtype: DT_INT32
        shape: (-1, 500)
        name: segment_ids:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['logits'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 56)
        name: loss/BiasAdd:0
  Method name is: tensorflow/serving/predict

运行转换脚本后，即可将SaveModel模型转换为FrozenGraph模型，此时通过Netron软件可查看模型详细的网络结构信息。

2. Tensorflow FrozenGraph转换为OM模型

可使用Ascend提供的OMG或ATC工具，将Tensorflow FrozenGraph转换为OM模型。

2.1 OMG模型转换

OMG(Offline Model Generator：离线模型生成器)支持运行在64位的Linux平台上，可在转换过程中进行算子调度优化、量化压缩、内存优化等，其使用方法可参考模型转换示例、OMG模型转换工具使用指导，支持的Tensorflow算子列表可参考Tensorflow算子边界。

如下所示，是Tensorflow BERT模型转换命令，因官方文档较为详细，本文不再详细介绍其使用方法。但需要注意的是，OMG支持的算子列表没有ATC工具丰富，因此建议使用ATC工具进行模型转换。

./omg --model ~/frozen_graph.pb --framework 3 --output ~/model --input_shape “input_ids:1,500;input_mask:1,500;segment_ids:1,500” --out_nodes “loss/BiasAdd”

2.2 ATC模型转换

ATC(Ascend Tensor Compiler：昇腾张量编译器)支持x86_64、aarch64架构，支持CentOS、EulerOS系统，可在转换过程中进行算子调度优化、量化压缩、内存优化等，支持的算子列表可参考Tensorflow框架算子规格。

使用ATC工具之前，需要搭建CANN(Compute Architecture for Neural Networks)环境，安装Ascend-cann-toolkit工具。详细安装教程可参考官方文档，因安装过程较为繁琐，可根据自身需求执行相应步骤，不必安装所有框架、插件。软件安装完成后，执行source /usr/local/Ascend/ascend-toolkit/set_env.sh命令来配置环境变量。执行atc -h命令，显示相关使用信息后，则说明atc工具配置成功。

ATC详细使用指南可参考快速入门，如下所示，是Tensorflow BERT模型转换命令，输出ATC run success显示模型转换成功。若转换过程中出现错误信息，可参考错误码参考。

atc --model=~/frozen_graph.pb --input_shape=“input_ids:1,500;input_mask:1,500;segment_ids:1,500” --framework=3 --output=model --soc_version=Ascend310

模型转换成功后，即可得到OM模型，可部署至Ascend机器上，如Ascend 310、Ascend 710，利用Ascend芯片来加速模型推理。

3. 其他问题

3.1 Pytorch模型如何转换为OM模型?

pytorch保存的模型，可先转换为ONNX模型，再利用OMG、ATC工具进行转换。

3.2 是否可将Tensorflow SaveModel转换为ONNX格式，再利用OMG、ATC转换为OM格式?

可以，可使用tf2onnx工具，将Tensorflow SaveModel转换为ONNX格式。但对于OMG、ATC工具而言，ONNX支持的算子没有Tensorflow算子多，因此建议转换为Tensorflow FrozenGraph。

3.2 转换过程中遇到不支持的算子如何解决？

针对不支持的算子，可选择的方法包括：1. 网络模型结构中去除相应算子，可参考Tensorflow: delete nodes from graph ；2. 不支持的算子，利用支持的算子进行替换实现。