ckpt轉(zhuǎn)換成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()

假設已經(jīng)生成了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模型轉(zhuǎn)換

tf serving是一款靈活的高性能機器學習服務系統(tǒng)，專為生產(chǎn)環(huán)境而設計。通過它可以輕松部署新算法和實驗，同時保持服務框架和API不變。它提供了與tensorflow模型的即是可用集成，但很容易擴展以便服務其他類型的模型和數(shù)據(jù)。

tf serving的安裝過程這里不多說，大家可以百度。

此處主要介紹tensorflow模型在docker中轉(zhuǎn)換時的修改內(nèi)容。

修改inception_saved_model.py文件中的內(nèi)容，主要包括：image_size，NUM_CLASSES，SYNSET_FILE，METADATA_FILE變量的內(nèi)容，必要時修改model_version，NUM_TOP_CLASSES。

修改inception_model.py文件中的內(nèi)容，包括從nets文件夾中導入所需網(wǎng)絡的信息，修改inference函數(shù)中對應的網(wǎng)絡名稱。

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環(huán)境時，需要添加環(huán)境及bazel編譯的配置項

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

ps，關(guān)于gpu的設置如下：

export CUDA_VISIBLE_DEVICES='0,1'  #shell環(huán)境
 
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" #python環(huán)境

以上這篇tensorflow轉(zhuǎn)換ckpt為savermodel模型的實現(xiàn)就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。

最新評論