Tensorflow pb模型多次加载推理的方法

       使用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模型进行加载，上述代码亦可稍微修改即使用。