如何在Spark中使用ONNX模型执行推理及实现方案合理性问询

你的实现思路是完全可行的，但在1亿条的大规模推理场景下，存在不少可优化的空间，先拆解原有方案的优劣势：

• 优势：
  • 逻辑直观，开发成本低，不需要额外引入推理服务、消息队列等组件，小批量验证可以快速跑通
  • ONNX格式相比原生transformers的PyTorch/TensorFlow模型，推理速度有30%~50%的提升，本身已经做了一层性能优化
• 劣势：
  • 如果用普通Spark Python UDF实现，JVM和Python进程之间的序列化、跨进程通信开销极高，这部分开销甚至会超过模型推理本身的耗时
  • 如果没有做模型单例加载，每个Executor的Task都会单独加载一次ONNX模型，会大量冗余占用内存/显存，很容易出现OOM问题，也会拉长任务的初始化时间

按改动成本从低到高，有几个梯度的优化方案可以选择：

方案1：最小改动优化（原有架构不变，性能提升3~10倍） #

只需要调整UDF的实现逻辑即可，不需要改整体流程：

• 把普通Python UDF替换为Pandas向量化UDF，按批次处理数据，每次传入128~1024条文本（可根据显存/内存大小调整），大幅降低JVM和Python的跨进程通信开销
• 用单例模式/lru_cache装饰器封装ONNX模型的加载逻辑，保证每个Executor进程仅加载一次模型，同进程下的所有Task共享模型实例，避免冗余加载
• 推理侧根据集群硬件选择对应的onnxruntime版本：CPU集群用onnxruntime-openvino（比基础版CPU推理快20%左右），GPU集群用onnxruntime-gpu，最大化推理速度

方案2：高性能CPU方案（性能再提升2~3倍） #

如果对任务耗时要求更高，可改用Java/Scala实现UDF：

• 直接调用ONNX Runtime的Java API加载模型，完全避开Python和JVM的跨进程通信开销
• 注意需要对齐Java侧的文本分词逻辑和你训练时transformers的分词逻辑，避免出现推理精度下降的问题
• 适合无GPU资源、完全依赖CPU集群做推理的场景

方案3：极致性能方案（性能提升10倍以上） #

如果集群有GPU资源，可接入Spark RAPIDS组件：

• 整个推理流水线全链路走GPU执行，包括文本分词、批处理、ONNX模型推理全部在GPU侧完成，不需要CPU和GPU之间频繁拷贝数据
• 1亿条短文本推理在中等规模GPU集群上可在数小时内跑完，远快于CPU集群的执行效率

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