tensorflow-ckpt转savemode记录

场景

我有一个tf1.X训练输出的ckpt格式的模型，需要转换为onnx格式。
直接用tf2onnx转onnx有问题，所以使用savemode作为中间格式，先将ckpt转为savemode，再将savemode转为onnx。

转为savemode的命令

importdocker  tensorflow.compat.v1 as tf
from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
#import tensorflow as tf

import os
tf.disable_eager_execution()

def read_graph_from_ckpt(ckpt_path, out_path ):     
    # 从meta文件加载网络结构
    saver = tf.train.import_meta_graph(ckpt_path+'.meta',clear_devices=True)
    graph = tf.get_default_graph()
    with tf.Session( graph=graph) as sess:
        sess.run(tf.global_variables_initializer()) 
        # 从ckpt加载参数
        saver.restore(sess, ckpt_path) 
        input_tf =graph.get_tensor_by_name('image_batch:0') 
        output_tf =graph.get_tensor_by_name('Logits_out/output:0') 
        
        #输入输出签名
        signature = predict_signature_def(inputs={'image_batch': input_tf},
                                  outputs={'Logits_out/output': output_tf})

        builder = tf.saved_model.builder.SavedModelBuilder(out_path)
        builder.add_meta_graph_and_variables(sess, ['serve'], strip_default_attrs=True, signature_def_map={'predict': signature})
        builder.save()
     
read_graph_from_ckpt('./tensorflow_landmark/model-tmp.ckpt', './saved_model')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

转为onnx的命令

python3 -m tf2onnx.convert --saved-model ./saved_model --opset 10 --output tf_face_landmark.onnx
1

遇到的错误

MetaGraphDef associated with tags ‘serve’ could not be found in SavedModel. To inspect available tag-sets in the SavedModel, please use the SavedModel CLI: saved_model_cli

调用tf2onnx.convert时我没指定–tag参数，所以默认的tag是“serve”。我原始的savemode转换代码没里也没加tag，所以就匹配不到。
解决方法在上面代码里，就是在add_meta_graph_and_variables第二个参数里指定“serve”为tag.

No signatures found in model. Try --concrete_function instead.

没有输入输出的签名，也就是没指定输入输出的节点。
解决方法就在上面代码里，就是增加signature_def_map这个参数。

吐槽下tensorflow

我在tf2.0里好用的代码，在tf2.2就出错了。兼容性真实一言难尽。
搞个转换onnx的活,tf2.X的环境下怎么都不行，我后来安装了一个py36+tf1.15+tf2onnx的环境才走通全流程。

参考资料

将tensorflow 1.x & 2.x转化成onnx文件（以arcface-tf2人脸识别模型为例）