美文网首页程序员的日常嵌入式开发深度学习
tensorRt加速tensorflow模型推理(incepti

tensorRt加速tensorflow模型推理(incepti

作者: 鱼香土豆丝 | 来源:发表于2018-12-08 17:37 被阅读197次

    摘要

    在一个人工智能大爆发的时代,一个企业不来点人工智能都不好意思说自己是科技企业。随着各公司在各自领域数据量的积累,以及深度学习的强拟合特点,各个公司都会训练出属于自己的模型,那么问题就来了,你有模型,我也有模型,那还比什么?对,就是速度,谁的速度快,谁就厉害。

    引言

    tensorflow作为谷歌开源的深度学习神器,被学界以及业界广泛的推广,本文就不做过多的介绍,如有需要可以参考如下链接

    人们往往在模型迭代指定次数后,保存一次checkpoint。在模型达到指定要求的精度后就可以上线应用了。NVIDIA tensorrt就是一种专门用于推理的引擎(不可训练),它主要是用于在生产环境中部署深度学习应用程序。以此来提供最大的吞吐量和效率。他主要是获取训练好的模型,这些模型通常是32或者16位的,tensorrt降低在这些精度(int8)来优化网络。具体的tensorrt介绍就不做过多的介绍。

    加速步骤(重点)

    训练好的inceptionV3模型保存的格式是一大堆的ckpt文件:

    在这里插入图片描述
    首先需要将这些文件转化为.pb的文件,也就是做模型的持久化,我们只需要模型做前向的推理,并得到预测值即可.这个步骤主要分为2步
    • 保存图
    • 把图和参数固定在一起

    这一步具体如何操作可以参考这篇博客
    接下来就需要构造tenssorrt的推理引擎了(以fp32的为例,int8的话还需要拿典型图片构造校准表)

    1.导入必要的库文件

    import uff
    import tensorrt as trt
    import pycuda.driver as cuda 
    import pycuda.autoinit
    from tensorrt.parsers import uffparser
    

    uff是将刚才的pb转化为引擎支持的uff文件,该文件可以序列化,也可以直接当作流传过去。
    trt则是用于加速推理的tensorrt
    pycyda则是用于显卡cuda编程的
    uffparser 解析uff模型的
    2.设置一些常量

    MODEL_DIR = './model_seg/model.pb'
    CHANNEL = 3
    HEIGHT = 299
    WIDTH = 299
    ENGINE_PATH = './model_seg/model_.pb.plan'
    INPUT_NODE = 'input'
    OUTPUT_NODE = 'InceptionV3/Logits/SpatialSqueeze'
    INPUT_SIZE = [CHANNEL, HEIGHT ,WIDTH] 
    MAX_BATCH_SIZE = 1 
    MAX_WORKSPACE = 1<<30
    

    MODEL_DIR 刚转的pb模型地址
    CHANNEL 图片的通道RGB 3 通道
    HEIGHT 图片的高
    WIDTH 图片的宽
    ENGINE_PATH 等会保存引擎的地址
    INPUT_NODE 模型的输入节点
    OUTPUT_NODE 模型的输出节点
    INPUT_SIZE 输入图片的大小,注意输入的是 CHANNEL, HEIGHT ,WIDTH
    MAX_BATCH_SIZE 每次输入几张图片吧
    MAX_WORKSPACE 显存的大小1<<30也就是1GB的大小,不信你试试
    3.pb转uff 并解析模型

    G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
    uff_model = uff.from_tensorflow_frozen_model(FROZEN_GDEF_PATH, [OUTPUT_NODE])
    parser = uffparser.create_uff_parser()
    parser.register_input(INPUT_NODE, INPUT_SIZE, 0)
    parser.register_output(OUTPUT_NODE)
    

    4.构建引擎

    engine = trt.utils.uff_to_trt_engine(G_LOGGER,uff_model,
                         parser,MAX_BATCH_SIZE,
                         MAX_WORKSPACE,datatype=trt.infer.DataType.FLOAT)
    

    5.保存模型

    trt.utils.cwrite_engine_to_file('./checkpoint/model_.pb.plan',engine.serialize())
    

    6.调用引擎实现推理

    def infer32(context, input_img, batch_size):
        engine = context.get_engine()
        assert(engine.get_nb_bindings() == 2)
        start = time.time()
        dims = engine.get_binding_dimensions(1).to_DimsCHW()
        elt_count = dims.C() * dims.H() * dims.W() * batch_size
        input_img = input_img.astype(np.float32)
        output = cuda.pagelocked_empty(elt_count, dtype=np.float32)
        d_input = cuda.mem_alloc(batch_size * input_img.size * input_img.dtype.itemsize)
        d_output = cuda.mem_alloc(batch_size * output.size * output.dtype.itemsize)
        bindings = [int(d_input), int(d_output)]
        stream = cuda.Stream()
        cuda.memcpy_htod_async(d_input, input_img, stream)
        context.enqueue(batch_size, bindings, stream.handle, None)
        cuda.memcpy_dtoh_async(output, d_output, stream)
        end = time.time()
        return output
    

    是不是很好懂?

    总结

    通过上述步骤可以将之前的tensorflow模型转化为rt的引擎,极大的提高推理速度,但是在flask中使用tensorrt的时候由于flask的上下文管理方式容易在成引擎找不到设备上下文,这一点可以用多进程的方式得到解决,在StackOverflow上有人提问这个问题,但是解决方案就是每推理一次都初始化一次cuda这样会影响整个应用的速度。只有单独的开一个进程供cuda进行初始化,这样才不会影响速度,也不会造成上下文冲突,具体的解决方法下一篇博客进行介绍

    相关文章

      网友评论

        本文标题:tensorRt加速tensorflow模型推理(incepti

        本文链接:https://www.haomeiwen.com/subject/wloqhqtx.html