Tensorflow PB模型转换为OM模型

举报
谓之小一 发表于 2022/04/12 12:27:14 2022/04/12
【摘要】 Tensorflow训练出的PB模型,若想在Ascend 310芯片上运行,需转换成OM模型,其主要转换流程包括:1. Tensorflow SaveModel转换为Tensorflow FrozenGraph;2. 利用OMG或ATC工具,将Tensorflow FrozenGraph模型转换为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. 不支持的算子,利用支持的算子进行替换实现。

【版权声明】本文为华为云社区用户原创内容,转载时必须标注文章的来源(华为云社区)、文章链接、文章作者等基本信息, 否则作者和本社区有权追究责任。如果您发现本社区中有涉嫌抄袭的内容,欢迎发送邮件进行举报,并提供相关证据,一经查实,本社区将立刻删除涉嫌侵权内容,举报邮箱: cloudbbs@huaweicloud.com
  • 点赞
  • 收藏
  • 关注作者

评论(0

0/1000
抱歉,系统识别当前为高风险访问,暂不支持该操作

全部回复

上滑加载中

设置昵称

在此一键设置昵称,即可参与社区互动!

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。

*长度不超过10个汉字或20个英文字符,设置后3个月内不可修改。