onnxruntime-gpu 预热速度优化

2024-04-23 09:54:33 浏览数 (1)

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

问题描述

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

代码语言:text复制
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

代码语言:text复制
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 就不一定

需要改的话:

代码语言:text复制
providers = [("CUDAExecutionProvider", {"cudnn_conv_use_max_workspace": '1'})]
io_binding

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

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

推理时:

代码语言:text复制
session.run_with_iobinding(binding)

在此之前需要创建 binding:

代码语言:text复制
binding = session.io_binding()

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

代码语言:text复制
# 输入 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(

)
Convolution Input Padding

卷积被转换成大矩阵乘法时,可以选择 N, C, D, 1 or N, C, 1, D 两种 pad 方式,结果相同,但由于会选择不同的卷积算法,导致性能可能不太一样。

特别是像 A100 这种显卡。

设置方式:

代码语言:text复制
providers = [("CUDAExecutionProvider", {"cudnn_conv1d_pad_to_nc1d": '1'})]

可以设置 0 和设置 1 都尝试一下,看看哪个更快。

参考资料

  • https://zhuanlan.zhihu.com/p/686755347

文章链接:

https://cloud.tencent.com/developer/article/2411347

0 人点赞