NVIDIA - TensorRT | onnxruntime

在 ONNX Runtime 中，TensorRT 执行引擎利用了 NVIDIA 的 TensorRT 深度学习推理引擎，从而提升了基于该引擎的 ONNX 模型在 NVIDIA GPU 上的运行效率。微软与 NVIDIA 密切合作，将 TensorRT 执行引擎成功集成到了 ONNX Runtime 中。

要使用 TensorRT 执行提供程序，必须在实例化InferenceSession时明确注册 TensorRT 执行提供程序。另外，建议也注册CUDAExecutionProvider，这样 Onnx Runtime 就能将那些 TensorRT 不支持的节点分配给 CUDA 执行提供程序来处理。

import onnxruntime as ort # set providers to ['TensorrtExecutionProvider', 'CUDAExecutionProvider'] with TensorrtExecutionProvider having the higher priority. sess = ort.InferenceSession('model.onnx', providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider'])

以下是设置 TensorRT EP 会话选项时的各种示例和场景：

import onnxruntime as ort model_path = '<path to model>' # note: for bool type options in python API, set them as False/True providers = [ ('TensorrtExecutionProvider', { 'device_id': 0, # Select GPU to execute 'trt_max_workspace_size': 2147483648, # Set GPU memory usage limit 'trt_fp16_enable': True, # Enable FP16 precision for faster inference }), ('CUDAExecutionProvider', { 'device_id': 0, 'arena_extend_strategy': 'kNextPowerOfTwo', 'gpu_mem_limit': 2 * 1024 * 1024 * 1024, 'cudnn_conv_algo_search': 'EXHAUSTIVE', 'do_copy_in_default_stream': True, }) ] sess_opt = ort.SessionOptions() sess = ort.InferenceSession(model_path, sess_options=sess_opt, providers=providers)