ONNX Runtime 创建推理会话

self.ort_session = ort.InferenceSession( onnx_model, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider'] )

加载一个.onnx模型文件，并根据当前设备情况选择 GPU 或 CPU 执行模型推理。

session = ort.InferenceSession( 模型路径, 执行设备 )

ONNX Runtime 是微软开发的推理框架。

ort.InferenceSession()

作用：

创建模型推理环境

model.onnx

onnx_model="GAT-RAE.onnx"

计算图 节点: Conv Linear Relu MatMul Attention 参数: W b 输入: x 输出: y

分析硬件：

CPU? GPU? CUDA? TensorRT?

选择执行器：

例如：

GPU:

CUDAExecutionProvider

CPU:

CPUExecutionProvider

onnx_model = "./model/gat_rae.onnx"

等价：

ort.InferenceSession( "./model/gat_rae.onnx" )

也可以：

onnx_model = bytes

直接加载二进制模型。

providers是什么？

ONNX Runtime 支持多个执行后端。

例如：

CPU

CPUExecutionProvider

使用：

Intel CPU AMD CPU

CUDA GPU

CUDAExecutionProvider

使用：

NVIDIA GPU

CUDAExecutionProvider

TensorRT

TensorrtExecutionProvider

用于：

NVIDIA推理优化

速度更快。

if ort.get_device() == 'GPU': providers=[ 'CUDAExecutionProvider', 'CPUExecutionProvider' ] else: providers=[ 'CPUExecutionProvider' ]

ort.get_device()

ort.get_device()

作用：

查看 ONNX Runtime 当前设备。

CPU

或者：

GPU