使用TensorRT加速Pytorch模型(简单实用)

目录

前言

TensorRT简介和安装

Pytorch模型转TensorRT

实验


 

前言

模型推理加速是调参师们都不得不面对的任务,加速一个模型有这么几个方向:使用混合精度,模型剪枝压缩,FPGA/ASIC硬件加速等。大部分加速方法,都需要大动干戈。而木盏直接推荐你两个简单实用的加速方法:1. 半精度(单精度); 2. TensorRT; 这两种方法的加速效果可以叠加。

关于第一种方法,可以看我的另一篇文章《【半精度】Pytorch模型加速和减少显存》,戳一戳戳不了吃亏。本文介绍TensorRT加速Pytorch模型。


TensorRT简介和安装

TensorRT是Nvidia官方给的C++推理加速工具,如同OpenVINO之于Intel。支持诸多的AI框架,如Tensorflow,Pytorch,Caffe,MXNet等。此外还对某些热门框架有特别的照顾,比如针对PyTorch有直接转换的工具torch2trt(咱们一会儿说)。

链接0:https://developer.nvidia.com/tensorrt

链接1:https://github.com/NVIDIA/TensorRT

链接2:https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html

链接3:https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/

TensorRT的主要原理是在GPU的硬件部署上对模型进行推理优化,合理分配硬件资源使得计算并行度更高。你可以根据你要的版本,下载tensorRT:

https://developer.nvidia.com/compute/machine-learning/tensorrt/secure/6.0/GA_6.0.1.5/tars/TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz

上面这串链接需要先登录nvidia账号,薛微有一点麻烦。一劳永逸嘛,不寒掺。

解压之:

tar -xvzf TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz

export之:

export TRT_RELEASE=`pwd`/TensorRT-6.0.1.5
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_RELEASE/lib

 安装python模块:

cd  TensorRT-6.0.1.5/python/
pip install tensorrt-6.0.1.5-cp37-none-linux_x86_64.whl

完毕~ 


Pytorch模型转TensorRT

这里用到的库是torch2trt,安装这个库要保证TensorRT已经安装好,而且cudatoolkit是对应PyTorch版本的(之前debug了好久发现的问题,希望你们不会遇到)。

看链接:https://github.com/NVIDIA-AI-IOT/torch2trt

本来点进去按照人家的README.md操作就行了。或者,看看我的踩坑经验。安装:

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

这是原demo:

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])

这里首先把pytorch模型加载到CUDA,然后定义好输入的样例x(这里主要用来指定输入的shape,用ones, zeros都可以)。model_trt就是转成功的TensorRT模型,你运行上面代码没报错就证明你转tensorRT成功了。

这里有一个小坑就是,原模型和tensorRT模型可能占2份GPU内存(额,也可能是我多虑,没做进一步实验)。那就可以先把tensorRT模型保存下来,下次推理的时候直接加载tensorRT模型就好:

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'))

推理的用法跟原来pytorch的用法一样:

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

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

实验

我针对自己训的CNN网络进行tensorRT加速,得到以下结果:

 单帧推理
baseline16.7ms
baseline + tensorRT8ms

在accuracy上,差异小于千分之一。快了1倍左右,所以TensorRT加速效果还可以

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 精致技术 设计师:CSDN官方博客 返回首页