tensorRt加速tensorflow模型推理(inception V3为例)

摘要

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

引言

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进行初始化,这样才不会影响速度,也不会造成上下文冲突,具体的解决方法下一篇博客进行介绍

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 216,193评论 6 498
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 92,306评论 3 392
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 162,130评论 0 353
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 58,110评论 1 292
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 67,118评论 6 388
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 51,085评论 1 295
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 40,007评论 3 417
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,844评论 0 273
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 45,283评论 1 310
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,508评论 2 332
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,667评论 1 348
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 35,395评论 5 343
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,985评论 3 325
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,630评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,797评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,653评论 2 368
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,553评论 2 352

推荐阅读更多精彩内容