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. 不支持的算子,利用支持的算子进行替换实现。
- 点赞
- 收藏
- 关注作者
评论(0)