TensorFlow模型预加载/一次加载多次推理

前言:深度学习算法的研发需要经过模型构建、训练、测试以及部署,其中在部署过程中的一个基本要求就是保证推理速度,而模型的加载会占用较多的时间,因此模型推理服务在调用前需要对模型进行预加载,即实现一次加载多次推理。

实现方法:在with语句下加载模型时,调用图(tf.Graph)与会话(tf.Session)的as_default()方法,将当前图与会话设置为默认图与默认会话,以保证在with语句外也可以执行当前会话,即实现一次加载多次执行。

具体实现模板:

# 一次推理多次调用

import tensorflow as tf

# 模型推理服务类
class ModelInferServer:
    def __init__(self):
        self.load_model()
    
    def load_model(self):
        # 创建tf.Graph
        self.graph = tf.Graph()

        # 创建tf.Session,创建时指定会话所对应的当前图self.graph
        self.sess = tf.Session(self.graph)
        
        # 加载模型时图与会话均调用as_default方法
        with self.graph.as_default():
            with self.sess.as_default():
                # 加载模型相关代码--根据模型情况编写 
                ...

    def infer(input):
        # 模型推理时图与会话同样调用as_default方法
        # 注意:若推理时as_default的调用在只调用一个模型的情况下为非必须
        #      若有多个不同的模型在一个类中封装的话则必须使用
        with self.graph.as_default():
            with self.sess.as_default():
                # 模型推理相关代码--根据模型情况编写 
                ...

if __name__ == "__main__":
    mis = ModelInferServer()
    
    # 多次推理调用样例

    # 第一次推理
    input1 = ...
    mis.infer(input1)
    # 第二次推理
    input2 = ...
    mis.infer(input2)
        

补充内容:若在模型推理服务类中仅仅封装了一个单一模型,则在模型推理时图与会话的as_default方法调用为非必须项,即以上代码块的第27、28行代码可以省略,相对地,若模型推理服务类中封装了多个不同模型,当前模型推理服务需要多个模型推理方能完成,类中会有多个图(如:self.graph1、self.graph2、...)与多个会话(如:self.sess1、self.sess2、...),同样会有多个推理方法(如:self.model1_infer、self.model2_infer、...),此时在调用具体的模型进行推理时就必须指定具体的图与会话为默认图与默认会话(如:self.graph1.as_default()、self.sess1.as_default())。

你可能感兴趣的:(python,tensorflow,tensorflow,人工智能,python)