tensorflow使用python对pb模型做预估

tensorflow中ckpt模型转成pb模型的代码：参考链接https://blog.youkuaiyun.com/dulingtingzi/article/details/90790282

但是为了使大家更容易明白，因为有些变量需要统一，这里针对下面的使用pb模型进行预估的代码，粘贴一下ckpt转pb模型：

import tensorflow as tf
import os
from tensorflow.contrib.layers import  flatten
from tensorflow.python.framework import graph_util
import tensorflow.contrib.slim as slim
import numpy as np

growth_rate = 6
depth = 50
compression = 0.5
weight_decay = 0.0001
nb_blocks = int((depth - 4) / 6)

def dense_net(img_input, num_classes, nb_blocks, growth_rate, weight_decay, compression, flag):
   ##自定义densenet代码
    return densenet(growth_rate,img_input,num_classes,weight_decay,nb_blocks,compression,flag)

def set_config():#设置GPU使用率# 控制使用率
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    # 假如有16GB的显存并使用其中的8GB:
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3)
    config = tf.ConfigProto(gpu_options=gpu_options)
    # session = tf.Session(config=config)
    return config
#下面是你自定义的模型
num_classes=2
#is_training = tf.placeholder(tf.bool, name='placeholder_is_training')
is_training = tf.constant(False, dtype=tf.bool)#下面的名字要和你一开始训练模型的时候是一致的
inputs = tf.placeholder(tf.float32, shape=[None,30, 280, 3], name='placeholder_x')
labels = tf.placeholder(tf.float32, shape=[None,num_classes], name='placeholder_y')
pred=dense_net(inputs, num_classes,nb_blocks, growth_rate,weight_decay,compression,is_training)

model_path="./version1/checkpoint/2_class.ckpt-1"#设置model的路径，因新版tensorflow会生成三个文件，只需写到数字前
cfg=set_config()

from tensorflow.python.saved_model import signature_constants, signature_def_utils, tag_constants, utils
save_path = './version1/model_pb/test'
with tf.Session(config=cfg) as sess:
    saver = tf.train.Saver()
    saver.restore(sess, model_path)
    print('ckpt loaded')
    #注意下面的inputs和outputs的名字要和后面的pb模型做预估保持一致
    model_signature = signature_def_utils.build_signature_def(inputs={"input": utils.build_tensor_info(inputs)},outputs={"pred": utils.build_tensor_info(pred)},method_name=signature_constants.PREDICT_METHOD_NAME)
    builder = tf.saved_model.builder.SavedModelBuilder(save_path)
    legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
    builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING],clear_devices=True,signature_def_map={signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:model_signature},legacy_init_op=legacy_init_op)
    builder.save()
    print('saved_model saved')

 

tensorflow使用pb模型做预估的代码：

import os
import math
import cv2
import numpy as np
import tensorflow as tf
from tensorflow.python.saved_model import signature_constants, signature_def_utils, tag_constants, utils
import matplotlib.pyplot as plt
from time import time

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

def preprocessing_crop_batch(images, height=30, width=280,depth=3):#按照batch去处理图像
    bs = len(images)
    GAUGE = height
    img_canvas = np.zeros([bs, height, width, depth], dtype=np.float32)
    #img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    for i,img in enumerate(images):
       #####预处理代码块
    return img_canvas

sess = tf.Session()
m = tf.saved_model.loader.load(sess, tags=[tag_constants.SERVING], export_dir='./version1/model_pb/test/')
graph = tf.get_default_graph()
signature = m.signature_def
signature_key = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
input_tensor_name0 = signature[signature_key].inputs['input'].name#与之前转pb的时候的输入输出名字保持一致
output_tensor_name = signature[signature_key].outputs['pred'].name
x0 = tf.get_default_graph().get_tensor_by_name(input_tensor_name0)
y0 = tf.get_default_graph().get_tensor_by_name(output_tensor_name)

img_path = './images'
imgs = os.listdir(img_path)
imgs = list(map(lambda x : os.path.join(img_path, x), imgs))

bs = 8
img_batch = []
for i in range(bs):
    img_batch.append(cv2.imread(imgs[i]))

t0 = time()
img_preprocess= preprocessing_crop_batch(img_batch)
ans = sess.run(y0, {x0 : img_preprocess})
ans = np.argmax(ans, axis=1)
#ans = list(map(lambda x : x.decode(), ans))
t1 = time()

#plt.imshow(img[:,:,::-1])
print(ans)
print('seconds per frame %.2f ms' % ((t1-t0) * 1000 / bs))