使用TensorRT加速Pytorch模型推理

2024年03月13日 由 alex 发表 392 0

TensorRT是 Nvidia 发布的一个机器学习框架,用于在其硬件上运行机器学习推理。TensorRT 经过高度优化,可在 NVIDIA GPU 上运行。这可能是目前运行模型最快的方法。


安装

我们应该下载与你的 CUDA 版本匹配的 TensorRT 格式。就我而言,我的 CUDA 版本是 12.1,因此我将使用相应的版本。


2


下载后,将其解压缩。之后,从 Extracted TensorRT 文件夹的 lib 文件夹中复制所有 .dll 文件。


1_-NvntQapFpVu6qUVcGaXCA


将它们粘贴到 CUDA 目录的 bin 文件夹中。


4


之后,转到Extracted TensorRT文件夹,然后在其中转到Python,并查找与你的Python版本相对应的文件。


5


就我而言,我使用的是 Python 3.10 版本。


6


你可以通过在 Python 环境中输入如上所示的 python — version 来检查你的 Python 版本。


然后对该文件进行 pip 安装,


7


然后在提取的 TensorRT 文件夹中找到 graphsurgeon 文件夹,并对其中的文件执行类似的 pip 安装,如下所示


8


9


接下来,对uff文件夹内的文件执行类似的 pip 安装,如下所示。


10


11


然后通过 pip install 重复相同的过程,但这次是针对onnx_grapghsurgeon文件夹内的文件。


12


13


一旦完成,我们就准备测试推理速度。为了进行测试,我们将使用 YOLOv8 预训练模型。


首先进行如下所示的 pip 安装,并将下载的 .pt 模型文件转换为 TensorRT 格式。


#!pip install ultralytics
from ultralytics import YOLO
model= YOLO('yolov8n.pt')
model.export(format='engine',device=0)


这将把 .pt 格式的 pytorch 模型导出为 .engine 格式。我们将比较同一数据集的 .pt 模型和 .engine 模型的性能,并测量它们的推理速度。


让我们使用下面的图片进行推理、


14


Pytorch .pt 模型性能


from ultralytics import YOLO
model= YOLO('yolov8n.pt',task='detect')
result= model.predict('japan-2014616_1280.jpg', save=True)


15


pytorch .pt 模型的推理时间为 99.5ms


TensorRT 模型


from ultralytics import YOLO
model= YOLO('yolov8n.engine',task='detect')
result= model.predict('japan-2014616_1280.jpg', save=True)


16


TensorRT 模型的推理时间为 2.0ms。现在,推理速度大幅降低。在未来几年中,企业应该转而使用 TensorRT 进行推理。感谢英伟达™(NVIDIA®)推出如此出色的软件包,缩短了推理时间,从而为企业节省了大量成本。


文章来源:https://levelup.gitconnected.com/tensorrt-makes-pytorch-inference-faster-6f56411599cd
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消