conda 环境安装onnx-runtime

1、ubuntu 系统中已安装cuda12.8

2、版本兼容

3、安装过程

# 1. 激活环境

conda create -n onnx_dev  python=3.10
conda activate onnx_dev

# 2. 安装 cuDNN(先搜索可用版本)
conda search cudnn -c conda-forge
# 选择版本
conda install cudnn==9.8.0.87 -c conda-forge
# 验证库文件是否存在
ls $CONDA_PREFIX/lib/libcudnn.so.9


# 3安装安装 TensorRT
pip install tensorrt==10.11.0.33 --extra-index-url https://pypi.nvidia.com
# 验证库文件是否存在
sudo find / -name "libnvinfer.so*" 2>/dev/null


# 4安装onnxruntime-gpu
onnxruntime-gpu==1.21.0

5 # 加入环境变量中
export LD_LIBRARY_PATH=/home/server/miniconda3/envs/ghm_onnx_dev/lib/python3.10/site-packages/tensorrt_libs:$LD_LIBRARY_PATH

3、验证

import onnxruntime as ort

print(ort.get_available_providers())

['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']

4、TensorrtExecutionProvider

TensorrtExecutionProvider(简称 TensorRT EP)是 ONNX Runtime 中的一个执行提供程序(Execution Provider),它的主要作用是将 ONNX 模型交给 NVIDIA TensorRT 引擎进行优化和推理,从而大幅提升在 NVIDIA GPU 上的运行性能

具体来说,它的作用和优势体现在以下几个方面:

  1. 深度硬件加速:它专门针对 NVIDIA RTX 或数据中心 GPU 进行了优化,能够充分利用 GPU 内部的 Tensor Core 以及 FP8、FP4 等底层架构优势,实现卓越的 AI 推理性能。
  2. 降低延迟与提升吞吐:相比于使用 DirectML 或普通的 CUDA 执行提供程序,TensorRT EP 能够显著降低推理延迟并提高吞吐量(例如在特定模型上吞吐量可提升 50%)。
  3. 即时编译与优化:它支持在最终用户设备上进行即时编译,自动对模型进行图优化和内核融合,生成高度优化的推理引擎,从而简化部署流程。
  4. 广泛的模型支持:它支持多种主流的 AI 模型架构,包括大语言模型(LLM)、扩散模型(Diffusion)以及传统的 CNN 等。

简而言之,如果你希望在 NVIDIA GPU 上获得极致的模型推理速度,在代码中指定 TensorrtExecutionProvider 是最佳选择。

6、示例代码

import onnxruntime as ort
import numpy as np

# 1. 配置 TensorRT 选项(可选,按需调整)
tensorrt_provider_options = {
    "trt_fp16_enable": True,          # 启用 FP16 混合精度,通常能显著加快推理速度
    "trt_engine_cache_enable": True,  # 开启引擎缓存
    "trt_engine_cache_path": "./trt_cache",  # 缓存编译后的引擎路径
}

# 2. 定义执行提供者列表(按优先级排序)
providers = [
    ("TensorrtExecutionProvider", tensorrt_provider_options),  # 优先使用 TensorRT
    ("CUDAExecutionProvider", {"device_id": 0}),               # 备选使用普通 CUDA
    "CPUExecutionProvider"                                     # 最终回退到 CPU
]

# 3. 创建推理会话
model_path = "your_model.onnx"  # 替换为你的模型路径
session = ort.InferenceSession(model_path, providers=providers)

# 4. 检查当前实际激活的提供者
print("当前激活的执行提供者:", session.get_providers()[0])

# 5. 准备输入数据并进行推理
input_name = session.get_inputs()[0].name
dummy_input = np.random.randn(1, 3, 224, 224).astype(np.float32)  # 示例输入

outputs = session.run(None, {input_name: dummy_input})
print("推理结果形状:", outputs[0].shape)

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值