如何优化动态输入下的ONNX模型推理速度？

我用ONNX Runtime创建会话的代码如下：

so = ort.SessionOptions()
so.inter_op_num_threads = 10
so.intra_op_num_threads = 10
session = ort.InferenceSession('example.onnx',
                               sess_options=so,
                               providers=['CUDAExecutionProvider'])

当输入尺寸固定（比如200）时，推理速度很快且稳定：

for i in tqdm(range(1000)):
    array = np.zeros((1, 200, 80), dtype=np.float32)
    embeddings = session.run(output_names=['embs'], input_feed={'feats': array})

但使用动态输入（尺寸在200-1000之间随机）时，前几百到几千次迭代速度极慢，之后才会优化到和固定输入相当的速度。尝试过批处理，但输入尺寸差异太大导致输出精度下降，求其他加速方法。

• 启用提前优化：在创建会话时开启ORT的全量优化选项，强制提前完成动态形状适配的优化，避免运行时动态编译开销。修改SessionOptions配置：

  so = ort.SessionOptions()
  so.inter_op_num_threads = 10
  so.intra_op_num_threads = 10
  # 启用所有级别的图优化，包含动态形状适配逻辑
  so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
  # 保存优化后的模型，后续加载可直接复用优化结果
  so.optimized_model_path = "optimized_example.onnx"
  session = ort.InferenceSession('example.onnx',
                                 sess_options=so,
                                 providers=['CUDAExecutionProvider'])
  

• 指定动态输入的形状范围：如果已知输入尺寸的上下限，在模型导出或会话创建阶段明确约束，让ORT提前生成适配该范围的优化代码，无需每次遇到新形状重新编译。
  以PyTorch导出ONNX为例，给动态维度设置明确范围：

  torch.onnx.export(model,
                    torch.randn(1, 200, 80),
                    "example.onnx",
                    dynamic_axes={"feats": {1: "seq_len"}},  # 标记第2维为动态维度
                    opset_version=17,
                    input_names=["feats"],
                    output_names=["embs"],
                    export_params=True,
                    do_constant_folding=True)
  

  或者给CUDA执行器配置提前搜索最优算法的策略：

  provider_options = [{
      "device_id": 0,
      "cudnn_conv_algo_search": "EXHAUSTIVE",  # 提前遍历最优卷积算法
      "enable_cuda_graph": True  # 启用CUDA Graph，配合形状缓存复用内核
  }]
  session = ort.InferenceSession('example.onnx',
                                 sess_options=so,
                                 providers=['CUDAExecutionProvider'],
                                 provider_options=provider_options)
  

• 预热高频输入形状：提前用业务中常见的输入形状跑几轮推理，让ORT提前编译好对应内核，消除正式运行时的初始延迟：

  # 预热业务中高频出现的形状
  warmup_shapes = [200, 400, 600, 800, 1000]
  for seq_len in warmup_shapes:
      array = np.zeros((1, seq_len, 80), dtype=np.float32)
      session.run(output_names=['embs'], input_feed={'feats': array})
  
  # 正式推理流程
  for i in tqdm(range(1000)):
      array = np.zeros((1, random.randint(200, 1000), 80), dtype=np.float32)
      embeddings = session.run(output_names=['embs'], input_feed={'feats': array})
  

• 调整形状缓存策略：ORT默认会缓存已编译的形状内核，可通过配置增大缓存容量，确保常用形状的内核都能被复用：

  so = ort.SessionOptions()
  so.inter_op_num_threads = 10
  so.intra_op_num_threads = 10
  so.enable_shape_inference = True
  so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
  # 增大形状缓存容量，避免频繁淘汰已编译内核
  so.set_session_config_entry("session.set_cache_capacity", "100")
  session = ort.InferenceSession('example.onnx',
                                 sess_options=so,
                                 providers=['CUDAExecutionProvider'])
  

内容的提问来源于stack exchange，提问作者Speсtra