本仓库包含使用 ONNX Runtime 和 TensorRT 优化 PyTorch 图像模型的代码，可实现高达 8 倍的推理加速。阅读完整博文请点击这里。

安装

创建并激活 conda 环境：

conda create -n supercharge_timm_tensorrt python=3.11
conda activate supercharge_timm_tensorrt

安装所需包：

pip install timm
pip install onnx
pip install onnxruntime-gpu==1.19.2
pip install cupy-cuda12x
pip install tensorrt==10.1.0 tensorrt-cu12==10.1.0 tensorrt-cu12-bindings==10.1.0 tensorrt-cu12-libs==10.1.0

安装 CUDA 依赖：

conda install -c nvidia cuda=12.2.2 cuda-tools=12.2.2 cuda-toolkit=12.2.2 cuda-version=12.2 cuda-command-line-tools=12.2.2 cuda-compiler=12.2.2 cuda-runtime=12.2.2

安装 cuDNN：

conda install cudnn==9.2.1.18

设置库路径：

export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="/home/dnth/mambaforge-pypy3/envs/supercharge_timm_tensorrt/lib/python3.11/site-packages/tensorrt_libs:$LD_LIBRARY_PATH"

运行代码

以下代码对应博文中的各个步骤。

PyTorch 延迟基准测试：

python 01_pytorch_latency_benchmark.py

了解更多请点击这里

将模型转换为 ONNX：

python 02_convert_to_onnx.py

了解更多请点击这里

ONNX Runtime CPU 推理：

python 03_onnx_cpu_inference.py

了解更多请点击这里

ONNX Runtime CUDA 推理：

python 04_onnx_cuda_inference.py

了解更多请点击这里

ONNX Runtime TensorRT 推理：

python 05_onnx_trt_inference.py

了解更多请点击这里

将预处理导出为 ONNX：

python 06_export_preprocessing_onnx.py

了解更多请点击这里

合并预处理和模型 ONNX：

python 07_onnx_compose_merge.py

了解更多请点击这里

对合并后的模型运行推理：

python 08_inference_merged_model.py

了解更多请点击这里

在视频上运行推理：

python 09_video_inference.py sample.mp4 output.mp4 --live 

<video controls autoplay src="https://cdn-uploads.huggingface.co/production/uploads/6195f404c07573b03c61702c/lOmu7KaqrihRDVcQVJDi0.mp4"></video>

使用摄像头作为输入源运行

python 09_video_inference.py --webcam --live