onnxruntime-gpu 预热速度优化

本文最后更新于：2024年5月7日下午

onnxruntime-gpu 在程序启动后第一次推断会消耗较大的系统资源，并且耗时更久，本文记录优化方法。

问题描述

在 Python 下 onnxruntime-gpu 加载 onnx 模型后，创建 seddion 进行数据推断，在第一次执行时会比之后执行耗时更久，需要资源更多。

session = onnxruntime.InferenceSession(str(model_path),  providers=[
                "CUDAExecutionProvider",
                "CPUExecutionProvider"
            ])
session.run(None, inputs)

解决方案

onnxruntime 的官方文档中有一些关于 Provider 的配置项说明：NVIDIA - CUDA | onnxruntime

其中 https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#cudnn_conv_algo_search

cudnn_conv_algo_search

The type of search done for cuDNN convolution algorithms.

Value	Description
EXHAUSTIVE (0)	expensive exhaustive benchmarking using cudnnFindConvolutionForwardAlgorithmEx
HEURISTIC (1)	lightweight heuristic based search using cudnnGetConvolutionForwardAlgorithm_v7
DEFAULT (2)	default algorithm using CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM

描述了Onnx 优化卷积操作的一个初始化搜索操作，在卷积多，而且 Onnx 需要接受多种可变尺寸输入时耗时严重，该选项默认为 EXHAUSTIVE，就是最耗时的那种。

因此如果遇到上述问题可以考虑尝试将该选项改为 DEFAULT

session = onnxruntime.InferenceSession(str(model_path), opts, providers=[
    ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
    "CPUExecutionProvider"
])

该选项优化在 Linux 下收益不太大，在 Windows 下可以将初始化预热时间从 500s 缩短到 70s。

其他性能调优

max_workspace

ORT 会使用 CuDNN 库来进行卷积计算，第一步是根据输入的 input shape, filter shape … 来决定使用哪一个卷积算法更好

需要预先分配 workspace，如果 workspace 不够大，有可能还执行不了最优的卷积算法

因此会想让 workspace 尽可能大，从而选择性能较好的卷积算法

1.14 以前的版本 cudnn_conv_use_max_workspace 这个 flag 默认是 0，意味着只会分配 32MB 出来，1.14 之后的版本默认是设置为 1，保证选择到最优的卷积算法，但有可能造成 peak memory usage 提高

官方说法是，fp16 模型，cudnn_conv_use_max_workspace 设置为 1 很重要，floatanddouble 就不一定

需要改的话：

1	`providers = [("CUDAExecutionProvider", {"cudnn_conv_use_max_workspace": '1'})]`

io_binding

可以减少一些数据拷贝（有时是设备间）的耗时。

如果要用这个，需要把 InferenceSession.run() 替换成 InferenceSession.run_with_iobinding()

推理时：

1	`session.run_with_iobinding(binding)`

在此之前需要创建 binding：

1	`binding = session.io_binding()`

把你需要的输入输出绑到 binding 上：

# 输入 X 来自 numpy array
io_binding.bind_cpu_input('X', X)

# 输入 X 来自 torch tensor
X_tensor = X.contiguous()
binding.bind_input(
    name='X',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(x_tensor.shape),
    buffer_ptr=x_tensor.data_ptr(),
)

# 让输出直接输出在一个 torch tensor 上
np_type = np.float32
DEVICE_NAME = 'cuda' if torch.cuda.is_available() else 'cpu'
DEVICE_INDEX = 0     # Replace this with the index of the device you want to run on
z_tensor = torch.empty(x_tensor.shape, dtype=torch_type, device=DEVICE).contiguous()
binding.bind_output(
    name='z',
    device_type=DEVICE_NAME,
    device_id=DEVICE_INDEX,
    element_type=np_type,
    shape=tuple(z_tensor.shape),
    buffer_ptr=z_tensor.data_ptr(),
)

# 让输出直接输出在 numpy array 上
binding.bind_output(

)