如何提升Keras网络单输入多次调用predict()的CPU推理速度

你遇到的性能缺口核心来源于推理框架调度开销远大于实际计算耗时，原生Keras的predict()接口每次调用都会产生图调度、数据拷贝等固定开销，对于你这种单步计算量极小的模型，这部分开销占比会超过90%，可以按以下优先级优化：

• 替换单次predict()调用
  不要对单个输入反复调用predict()，有两种优化方案：

  1. 批处理推理：把多次推理的输入攒成一个批次，一次性传入predict()，批次大小可设为16/32/64，充分利用CPU并行能力
  2. 实时场景下直接调用模型对象：用model(input_tensor, training=False)替代predict()，可直接跳过predict()的额外封装开销，同时要显式传入training=False关闭Dropout层的推理阶段随机计算，既提升精度也降低无用开销
  3. 用tf.function固化计算图：提前执行infer_func = tf.function(model, jit_compile=True)，后续直接调用infer_func(input_tensor)，开启XLA编译优化，可进一步压缩调度开销，预期单线程推理速度能提升10倍以上

• 模型结构与量化优化

  1. 把模型中的transpose操作移到预处理阶段执行，避免每次推理都重复执行张量转置操作
  2. 执行INT8量化：通过TensorFlow Lite转换器或者ONNX量化工具把模型转为INT8精度，对于你这种小模型，精度损失几乎可忽略，推理速度可再提升2~4倍
  3. 可尝试裁剪冗余卷积通道：你的模型卷积通道数有冗余，可通过剪枝工具进一步压缩参数量和计算量

• 更换推理运行时
  放弃原生TensorFlow/Keras推理，改用专门的CPU推理运行时：

  1. 转ONNX格式后用ONNX Runtime运行，默认开启CPU指令集优化，速度比原生TF高3~5倍
  2. 若使用Intel CPU，可转成OpenVINO格式运行，针对AVX2/AVX512指令集做深度优化，性能还能再提升20%~50%

• 运行时参数调整
  显式设置TensorFlow线程参数，避免默认线程调度开销：

  import tensorflow as tf
  # 物理核心数按实际CPU配置填写，比如4核就填4
  tf.config.threading.set_intra_op_parallelism_threads(物理核心数)
  tf.config.threading.set_inter_op_parallelism_threads(1)
  

  关闭Eager Execution模式，强制使用静态图执行推理，也可降低动态调度开销。

按以上方案优化后，基本可以达到你预期的每秒1000次左右的推理性能。

内容的提问来源于stack exchange，提问作者Rob