Chinaunix首页 | 论坛 | 博客
  • 博客访问: 3882834
  • 博文数量: 356
  • 博客积分: 10458
  • 博客等级: 上将
  • 技术积分: 4734
  • 用 户 组: 普通用户
  • 注册时间: 2008-03-24 14:59
文章分类

全部博文(356)

文章存档

2020年(17)

2019年(9)

2018年(26)

2017年(5)

2016年(11)

2015年(20)

2014年(2)

2013年(17)

2012年(15)

2011年(4)

2010年(7)

2009年(14)

2008年(209)

分类: Python/Ruby

2020-03-24 11:11:03

ckpt转换成SavedModel

convert_ckpt_to_savermodel.py

点击(此处)折叠或打开

  1. import tensorflow as tf
  2. import sys

  3. trained_checkpoint_prefix = sys.argv[1]
  4. export_dir = sys.argv[2]
  5. graph = tf.Graph()
  6. config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
  7. with tf.compat.v1.Session(graph=graph, config=config) as sess:
  8.     # Restore from checkpoint
  9.     loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
  10.     loader.restore(sess, trained_checkpoint_prefix)

  11.     # Export checkpoint to SavedModel
  12.     builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
  13.     builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)
  14.     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不存在的错误。

作者:帅得不敢出门
阅读(766) | 评论(0) | 转发(0) |
给主人留下些什么吧!~~