TensorRT是 Nvidia 发布的一个机器学习框架,用于在其硬件上运行机器学习推理。TensorRT 经过高度优化,可在 NVIDIA GPU 上运行。这可能是目前运行模型最快的方法。
安装
我们应该下载与你的 CUDA 版本匹配的 TensorRT 格式。就我而言,我的 CUDA 版本是 12.1,因此我将使用相应的版本。
下载后,将其解压缩。之后,从 Extracted TensorRT 文件夹的 lib 文件夹中复制所有 .dll 文件。
将它们粘贴到 CUDA 目录的 bin 文件夹中。
之后,转到Extracted TensorRT文件夹,然后在其中转到Python,并查找与你的Python版本相对应的文件。
就我而言,我使用的是 Python 3.10 版本。
你可以通过在 Python 环境中输入如上所示的 python — version 来检查你的 Python 版本。
然后对该文件进行 pip 安装,
然后在提取的 TensorRT 文件夹中找到 graphsurgeon 文件夹,并对其中的文件执行类似的 pip 安装,如下所示
接下来,对uff文件夹内的文件执行类似的 pip 安装,如下所示。
然后通过 pip install 重复相同的过程,但这次是针对onnx_grapghsurgeon文件夹内的文件。
一旦完成,我们就准备测试推理速度。为了进行测试,我们将使用 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 模型的性能,并测量它们的推理速度。
让我们使用下面的图片进行推理、
Pytorch .pt 模型性能
from ultralytics import YOLO
model= YOLO('yolov8n.pt',task='detect')
result= model.predict('japan-2014616_1280.jpg', save=True)
pytorch .pt 模型的推理时间为 99.5ms
TensorRT 模型
from ultralytics import YOLO
model= YOLO('yolov8n.engine',task='detect')
result= model.predict('japan-2014616_1280.jpg', save=True)
TensorRT 模型的推理时间为 2.0ms。现在,推理速度大幅降低。在未来几年中,企业应该转而使用 TensorRT 进行推理。感谢英伟达™(NVIDIA®)推出如此出色的软件包,缩短了推理时间,从而为企业节省了大量成本。