模型部署之 TensorRT 初步入门

TensorRT 是 NVIDIA 提出的用于统一模型部署的加速器,可以应用于 NVIDIA 自家设计的硬件平台比如 NVIDIA Tesla A100 显卡,JETSON Xavier 开发板等,它的输入可以是来自各个流行的训练框架,比如 Tensorflow, Pytorch 等训练得到的模型结果。

官网定义:

TensorRT is built on CUDA, NVIDIA’s parallel programming model, and enables you to optimize inference for all deep learning frameworks leveraging libraries, development tools and technologies in CUDA-X for artificial intelligence, autonomous machines, high-performance computing, and graphics.

TensorRT包括 推理优化(inference optimization) 和 runtime 两部分,类似于 MicroSoft 提出的 ONNX Runtime,但 ONNX Runtime 一般只能接收 ONNX 格式的模型,而TensorRT可以接受包括ONNX,Pytorch, Tensorflow等基本上所有框架的模型


TensorRT在对模型优化时主要进行了5个调整:

1. Layer and tensor fusion

kernel fusion

kernel fusion 的主要目的是提高GPU的利用效率,减少 kernel 的数目,因为每增加一个算子就会增加一份数据读写,而数据读写是相对比较耗时的,同时增加一个算子也会增加一次计算

  • 因此将可以融合的模块比如 conv-bn-relu三个模块就可以融合成一个模块这就可以减少数据读写和多次计算
  • 对于具有同一个输入,且模块内容相同的模块,但是输出不一样的,如上图左的 3 个 1x1 模块,就可以利用并行parallel进行计算,再输出到不同的节点。具体的实现方法后续跟进

When you have identical kernels which take the same input but just use different weights, you can combine the kernels by making a single kernel wider in a sense that is processes more of these operations in parallel. The output from these horizontally fused kernels will be automatically split up if they feed to different kernels further down the graph。

2. Precision Calibration

校准精度,由于这里 inference 过程只需要 forward,并不需要 backward, 因此就不需要 32位的浮点数来进行计算,因此可以合理的采用 fp16 或者 int8 来进行 forward, 这样可以是的模型存储空间更小,更低的内存占用和延迟



具体的实现方法后续跟进,引用如下:

TensorRT achieved this by using an automated parameter-free calibration step to change the weighs and activation tensors into lower precision using a representative input sample and this is done such that the model minimizes the accuracy loss.

3. Kernel Auto-tuning

对于同一个操作(卷积等)有很多不同的底层实现,TensorRT 可以根据你的参数 比如 batch-size, filter-size, input data size 等或者部署平台去选择最优的实现方法。


4. Dynamic Tensor Memory

dynamic tensor memory ensures that memory is allocated for each tensor only for the duration of its usage. This naturally reduces memory footprint and improves memory reuse.

5. Multi-Stream Execution

Multi-stream execution is essential when you scale the inference to multiple clients. This is achieved by allowing multiple input streams to use the same model in parallel on a single device


代码:

可以使用 TRTorch, torch2trt, 或者TF-TRT对模型进行转换
TRTorch, torch2trt

image.png

pytorch 举例:

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()

# create example data
x = torch.ones((1, 3, 224, 224)).cuda()

# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

y = model(x)
y_trt = model_trt(x)

# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

模型保存和加载:

torch.save(model_trt.state_dict(), 'alexnet_trt.pth')
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))

Webinar: Deploying Models with TensorRT

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

推荐阅读更多精彩内容