使用TensorFlow进行模型训练后,若需要对同一个pb模型多个对象调用推理时,此时可以通过创建类的方式将加载模型的步骤置于类中的初始化便能够通过创建多个类的对象进行pb模型的加载以及推理。
实际代码如下:
import tensorflow as tf
import numpy as np
import cv2
from PIL import Image
from tensorflow.python.platform import gfile
import os
# import keras.backend.tensorflow_backend as KTF
class pbInference(object):
def __init__(self, PATH_TO_CKPT, WIDTH_HEIGH):
pid = os.getpid()
print('------init pid-----', pid)
self.WIDTH_HEIGH = WIDTH_HEIGH
self.sess = self.config(gpu_num=gpu_num)
#self.sess = tf.Session(graph=g)
with gfile.FastGFile(PATH_TO_CKPT, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
self.sess.graph.as_default()
tf.import_graph_def(graph_def, name='') # 导入计算图
self.sess.run(tf.global_variables_initializer())
def config(self, gpu_num='0'):
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_num
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
# KTF.set_session(session)
return session
def load_image_into_numpy_array(self, image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape(
(im_height, im_width, 3)).astype(np.uint8)
def service(self, job_id, image_path):
image = Image.open(image_path)
image = image.convert('RGB')
image_np = self.load_image_into_numpy_array(image)
image_np = cv2.resize(image_np, (self.WIDTH_HEIGH, self.WIDTH_HEIGH))
image_np_expanded = np.expand_dims(image_np, axis=0)
input = self.sess.graph.get_tensor_by_name("input:0")
label = self.sess.graph.get_tensor_by_name("output:0")
output_dict = self.sess.run(label, feed_dict={input: image_np_expanded})
print('output_dict:',output_dict)
label = np.array(output_dict[0]).argmax()
score = output_dict[0][label]
return label,score
if __name__ == '__main__':
# t1 = time.time()
file = xxx.jpg'
PATH_TO_CKPT = './pb/freezn_graph.pb'
WIDTH_HEIGH = 224
gpu_num = '3'
g = pbInference(PATH_TO_CKPT=PATH_TO_CKPT, WIDTH_HEIGH=WIDTH_HEIGH)
label,score = g.service("1",file)
print('label:',label)
print('score:',score)
以上便是通过创建类的方式进行pb模型加载调用,若需要对多个pb模型进行加载,上述代码亦可稍微修改即使用。