【Pytorch】ONNX Runtime + TensorRT + PyCUDA 安装记录

Xenon_Xu

1249人浏览 · 2026-01-05 17:00:48

Xenon_Xu · 2026-01-05 17:00:48 发布

任务背景

计划在 CentOS 服务器部署 torch 模型推理，需要安装 ONNX Runtime + TensorRT + PyCUDA 依赖库。
所用显卡型号为 NVIDIA RTX 4090，CUDA Runtime Version 12.8

CUDA、PyTorch、cuDNN、onnx、onnx-runtime、TensorRT 均有严格的版本匹配要求，直接安装有较高概率 crash

由于网上安装流程、资料不完善，安装过程不断试错过于痛苦，
本文将对完整的安装思路、过程进行记录，并附上测试脚本，供大家参考。

具体步骤

Step 0：构建虚拟环境

本人使用 Anaconda 创建 python=3.11 的虚拟环境，具体过程省略。

Step 1：确定 CUDA、PyTorch、cuDNN 版本

nvidia-smi查看 CUDA 运行时版本：
nvidia-smi查看CUDA版本
本人所用 CUDA 版本 12.8，安装（更新）PyTorch 选取一个最保守的版本。
根据 https://pytorch.org/get-started/previous-versions/ 搜索 12.8，找到

# CUDA 12.8
pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128

得知 CUDA 12.8 的最保守 torch 版本为 2.7.0+cu128，安装（或更新）该版本 torch：

pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128

注：安装旧版本（例如 torch==2.0.1）torch 可以在 CUDA 12.8 上运行，但大概率导致后续 onnx-runtime、tensorRT 无法正常运行，故强烈建议遵循 PyTorch 官方推荐的版本号

进入虚拟环境，运行 python，查看 cuDNN 版本：

(virtual-env)[user@server pwd]$ python
Python 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
Ctrl click to launch VS Code Native REPL
>>> import torch
>>> torch.backends.cudnn.version()
90701

确定 cuDNN 版本为 9.07.01，即大版本号为9.x。

Step 2：确定 ONNX Runtime、TensorRT 版本号

根据 https://onnxruntime.ai/docs/execution-providers/TensorRT-ExecutionProvider.html，找到支持 CUDA 12.8 的最保守版本：
确认ONNX版本
确定 ONNX Runtime 版本为 1.21、TensorRT 版本为 10.8。

根据 https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html 再次检查 ONNX Runtime 与 CUDA、cuDNN 的版本兼容性：
ONNX-Runtime版本兼容

布什鸽门你的1.21.x呢？

onnx-runtime 1.20.x 已经要求 cuDNN 大版本为 9.x，torch 版本高于 2.4.0，推测 onnx-runtime 1.21.x 同理。
目前 torch 2.7.0+cu128、cuDNN 9.07.01 满足要求。若这里版本号出现不匹配，可以微调 torch、onnx-runtime 的版本号，直到全部匹配。

现在去确定 onnx 的版本，根据 https://onnxruntime.ai/docs/reference/compatibility.html：
onnx版本兼容

布什鸽门你的1.21.x呢？？

但我们拿到一个重要信息，即 onnx-runtime>1.20，对应的 IR version>=10、opset version>=21
根据 https://github.com/onnx/onnx/blob/main/docs/Versioning.md，找到不同 onnx 对应的 IR Version：
onnx对应ir版本
我们再次保守假设 onnx-runtime 1.21.x 可能只支持 IR version>=10、opset version>=21，即 onnx<=1.16.2。
但本人实际安装时，大胆往前推进了一个版本 onnx==1.17.0，假设 onnx-runtime 1.21.x 支持 IR version>=10、opset version>=22。

在 Step 3 安装后，我们可以进行测试：

# test onnx-runtime-gpu opset version
import onnx
from onnx import helper, TensorProto
import onnxruntime as ort

def build_model(opset):
    # 构建一个最简单的 Add 模型
    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])
    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])
    Z = helper.make_tensor_value_info("Z", TensorProto.FLOAT, [1])

    node = helper.make_node("Add", ["X", "Y"], ["Z"])
    graph = helper.make_graph([node], "test_graph", [X, Y], [Z])
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)])
    return model

def test_opset(opset):
    model = build_model(opset)
    path = f"test_opset_{opset}.onnx"
    onnx.save(model, path)

    try:
        sess = ort.InferenceSession(path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
        return True
    except Exception as e:
        return False

supported = []
for opset in range(10, 25):  # 测试 10~24
    ok = test_opset(opset)
    print(f"Opset {opset}: {'OK' if ok else 'NOT SUPPORTED'}")
    if ok:
        supported.append(opset)

print("\n最大支持 Opset =", max(supported))

onnx==1.17.0、onnx-runtime==1.21.0 测试结果为：

Opset 20: OK
Opset 21: OK
Opset 22: OK
Opset 23: NOT SUPPORTED
Opset 24: NOT SUPPORTED
最大支持 Opset = 22

说明正好猜对了哈哈~
(测试完了记得把 onnx 缓存删了)

如果 Step 3 安装验证后发现 opset version 不匹配，可以再回退 onnx 版本。
以上确定 onnx 版本为 1.17.0。

Step 3：安装 onnx、onnx-runtime-gpu、pycuda

执行 pip 安装（注意版本号）：

pip install onnxruntime-gpu==1.21 onnx==1.17 -i https://pypi.tuna.tsinghua.edu.cn/simple

可以配置镜像源（例如清华源 https://pypi.tuna.tsinghua.edu.cn/simple）加速安装。

PyCUDA 安装最简单，直接 pip 即可：

pip install pycuda

本人pycuda的安装结果如下：

Successfully installed mako-1.3.10 pycuda-2025.1.2 pytools-2025.2.5 siphash24-1.8

Step 4：安装 TensorRT

这里建议使用 NVIDIA 官方的 TAR 安装方法。
进入 https://developer.nvidia.com/tensorrt/download/10x，找到 10.8 GA，CUDA 12.8 的 TAR 文件：

需要注册 NVIDIA Developer 账号才能下载
更多 TensorRT 版本详见 https://developer.nvidia.com/tensorrt/download

将 tar 下载到服务器（或者本地下载、再上传至服务器），本人下载到服务器的路径为 ~/local/。

TAR 文件非常大（TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz 有 6 个GB）

然后解压至 ~/local/路径：

tar -xvf TensorRT-10.8.0.43.Linux.x86_64-gnu.cuda-12.8.tar.gz -C ~/local/

执行 pip 安装 whl 文件（注意文件名可能不同）：

pip install ~/local/TensorRT-10.8.0.43/python/tensorrt-10.8.0.43-cp311-none-linux_x86_64.whl

Step 5：测试

非常重要：启动 python 程序前，必须保证 `$LD_LIBRARY_PATH` 配置正确：

export LD_LIBRARY_PATH=$HOME/local/TensorRT-10.8.0.43/lib:$HOME/anaconda3/envs/你的环境名称/lib/python3.11/site-packages/nvidia/cudnn/lib:$LD_LIBRARY_PATH

$HOME/local/TensorRT-10.8.0.43/lib：即 TensorRT 的 so 运行库位置；

你的环境路径/lib/python3.11/site-packages/nvidia/cudnn/lib：即虚拟环境中 cuDNN 的 so 运行库位置。

可以写入 ~/.bashrc 实现永久配置。

测试脚本如下：

# verify_env.py
print("=== 环境验证开始 ===")

# 1. 验证 ONNX
try:
    import onnx
    from onnx import helper
    node = helper.make_node("Relu", inputs=["x"], outputs=["y"])
    print("[ONNX] 成功导入，能创建简单算子节点:", node.op_type)
except Exception as e:
    print("[ONNX] 出错:", e)

# 2. 验证 ONNXRuntime
try:
    import onnxruntime as ort
    sess = ort.InferenceSession(
        onnx.helper.make_model(
            onnx.helper.make_graph([helper.make_node("Relu", ["x"], ["y"])],
                                   "test_graph",
                                   [helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1])],
                                   [helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])]
            ),
            # opset_imports=[helper.make_opsetid("", 11)]
        ).SerializeToString(),
        # providers=["CPUExecutionProvider"]
        providers=["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
    )
    print("[ONNXRuntime] 成功创建推理 Session，Provider:", sess.get_providers())
except Exception as e:
    print("[ONNXRuntime] 出错:", e)

# 3. 验证 TensorRT
try:
    import tensorrt as trt
    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(TRT_LOGGER)
    print("[TensorRT] 成功导入，版本:", trt.__version__)
except Exception as e:
    print("[TensorRT] 出错:", e)

# 4. 验证 PyCUDA
try:
    import pycuda.driver as cuda
    import pycuda.autoinit
    print("[PyCUDA] 成功导入，GPU 名称:", cuda.Device(0).name())
except Exception as e:
    print("[PyCUDA] 出错:", e)

print("=== 环境验证结束 ===")

本人运行结果为：

=== 环境验证开始 ===
[ONNX] 成功导入，能创建简单算子节点: Relu
[ONNXRuntime] 成功创建推理 Session，Provider: ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
[TensorRT] 成功导入，版本: 10.8.0.43
[PyCUDA] 成功导入，GPU 名称: NVIDIA GeForce RTX 4090
=== 环境验证结束 ===

若没有任何警告、报错信息，且 onnx-runtime provider 包含 tensorrt、cuda，则安装成功。

若 TensorRT 出错：

说明找不到 TensorRT 运行库，请检查 LD_LIBRARY_PATH 是否在运行 python 程序前配置了 TensorRT 的 lib 路径；若 .so.x 显示的是其它数字（例如.so.8）说明装错了，回到 step 2、4 重装 TensorRT。
若 onnx runtime 出错：

说明找不到 cuDNN 运行库，请检查 LD_LIBRARY_PATH 是否在运行 python 程序前配置了 cuDNN 的 lib 路径；若 .so.x 显示的是其它数字（例如.so.8）说明装错了，回到 step 1 重装 PyTorch、cuDNN。
其它问题不一一列举了……