onnxruntime-gpu 使用和配置

最近把 F5-TTS 导出了 onnx，这个是用 GPU 的时候关于 onnxruntime-gpu 的配置。

pip install onnxruntime-gpu onnxruntime-tools -i https://pypi.tuna.tsinghua.edu.cn/simple

注意版本的对应：

检查是否支持的 GPU 环境:

import onnxruntime as ort
print(ort.__version__)  # 查看 onnxruntime 版本

# 检查 GPU 是否在执行提供程序中
providers = ort.get_all_providers()

# ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'MIGraphXExecutionProvider', 'ROCMExecutionProvider', 'OpenVINOExecutionProvider', 'DnnlExecutionProvider', 'TvmExecutionProvider', 'VitisAIExecutionProvider', 'QNNExecutionProvider', 'NnapiExecutionProvider', 'VSINPUExecutionProvider', 'JsExecutionProvider', 'CoreMLExecutionProvider', 'ArmNNExecutionProvider', 'ACLExecutionProvider', 'DmlExecutionProvider', 'RknpuExecutionProvider', 'WebNNExecutionProvider', 'WebGpuExecutionProvider', 'XnnpackExecutionProvider', 'CANNExecutionProvider', 'AzureExecutionProvider', 'CPUExecutionProvider']

if 'CUDAExecutionProvider' in providers:
    print("CUDA GPU support is available")
else:
    print("CUDA GPU support is not available")

指定 CUDAExecutionProvider：

import onnxruntime as ort
sess_options = ort.SessionOptions()
session = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"], sess_options=sess_options)

量化:

from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic("model.onnx", "model_quant.onnx", weight_type=QuantType.QInt8)

同时也可以简单测一下：

# 创建 ONNX Runtime 推理会话
session = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider"])

# 测试推理
input_data = {"input_name": ...}  # 替换为实际输入
outputs = session.run(None, input_data)
print(outputs)