tensorflow转换ckpt为savermodel模型

ckpt转换成SavedModel

convert_ckpt_to_savermodel.py

点击(此处)折叠或打开

1. import tensorflow as tf
   
2. import sys
   
3. 
   
4. trained_checkpoint_prefix = sys.argv[1]
   
5. export_dir = sys.argv[2]
   
6. graph = tf.Graph()
   
7. config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
   
8. with tf.compat.v1.Session(graph=graph, config=config) as sess:
   
9.     # Restore from checkpoint
   
10.     loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
    
11.     loader.restore(sess, trained_checkpoint_prefix)
    
12. 
    
13.     # Export checkpoint to SavedModel
    
14.     builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
    
15.     builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)
    
16.     builder.save()

假设已经生成了ckpt模型checkpoint   hello_model.data-00000-of-00001  hello_model.index  hello_model.meta

python ./convert_ckpt_to_savermodel.py  hello_model ./save

会在save目录下生成

save
├── saved_model.pb
└── variables
   ├── variables.data-00000-of-00001
   └── variables.index

[tf.saved_model.TRAINING, tf.saved_model.SERVING]  可根据需要修改此列表，比如如果模型中只有SERVING那么要改成
[tf.saved_model.SERVING] ,否则会提示 TRAINING不存在的错误。

作者：帅得不敢出门