tensorflow转换ckpt为savermodel模型的实现

(编辑:jimmy 日期: 2024/11/14 浏览:2)

ckpt转换成SavedModel

convert_ckpt_to_savermodel.py

import tensorflow as tf
import sys
 
trained_checkpoint_prefix = sys.argv[1]
export_dir = sys.argv[2]
graph = tf.Graph()
config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
with tf.compat.v1.Session(graph=graph, config=config) as sess:
  # Restore from checkpoint
  loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
  loader.restore(sess, trained_checkpoint_prefix)
 
  # Export checkpoint to SavedModel
  builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
  builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)
  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

补充知识:tensorflow serving模型转换

tf serving是一款灵活的高性能机器学习服务系统,专为生产环境而设计。通过它可以轻松部署新算法和实验,同时保持服务框架和API不变。它提供了与tensorflow模型的即是可用集成,但很容易扩展以便服务其他类型的模型和数据。

tf serving的安装过程这里不多说,大家可以百度。

此处主要介绍tensorflow模型在docker中转换时的修改内容。

修改inception_saved_model.py文件中的内容,主要包括:image_size,NUM_CLASSES,SYNSET_FILE,METADATA_FILE变量的内容,必要时修改model_version,NUM_TOP_CLASSES。

修改inception_model.py文件中的内容,包括从nets文件夹中导入所需网络的信息,修改inference函数中对应的网络名称。

from nets.inception_v1 import inception_v1, inception_v1_arg_scope

with slim.arg_scope(inception_v1_arg_scope()):
  logits, endpoints = inception_v1(
     images,
     dropout_keep_prob=0.8,
     num_classes=num_classes,
     is_training=for_training,
     scope=scope)

另,使用CUDA环境时,需要添加环境及bazel编译的配置项

export TF_NEED_CUDA=1
bazel build -c opt --config=cuda tf_models/slim:inception_saved_model

ps,关于gpu的设置如下:

export CUDA_VISIBLE_DEVICES='0,1'  #shell环境
 
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" #python环境

以上这篇tensorflow转换ckpt为savermodel模型的实现就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持。

一句话新闻
高通与谷歌联手!首款骁龙PC优化Chrome浏览器发布
高通和谷歌日前宣布,推出首次面向搭载骁龙的Windows PC的优化版Chrome浏览器。
在对骁龙X Elite参考设计的初步测试中,全新的Chrome浏览器在Speedometer 2.1基准测试中实现了显著的性能提升。
预计在2024年年中之前,搭载骁龙X Elite计算平台的PC将面世。该浏览器的提前问世,有助于骁龙PC问世就获得满血表现。
谷歌高级副总裁Hiroshi Lockheimer表示,此次与高通的合作将有助于确保Chrome用户在当前ARM兼容的PC上获得最佳的浏览体验。