SynapseML 实战：使用 ONNX 模型在 Spark 上进行分布式推理-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00538/article/details/148506382

SynapseML 实战：使用 ONNX 模型在 Spark 上进行分布式推理

SynapseML microsoft/SynapseML: 是一个开源的机器学习框架，用于构建和部署人工智能应用。它提供了丰富的机器学习算法和工具，可以帮助开发者快速构建 AI 应用。特点包括易于使用、高性能、支持多种机器学习算法等。项目地址: https://gitcode.com/gh_mirrors/sy/SynapseML

前言

在大规模机器学习场景中，模型训练完成后如何高效地进行分布式推理是一个关键问题。本文将介绍如何利用 SynapseML 框架，将训练好的 LightGBM 模型转换为 ONNX 格式，并在 Spark 集群上实现高效的分布式推理。

环境准备

在开始之前，我们需要确保环境中安装了必要的 Python 包：

%pip install --no-cache-dir lightgbm onnxmltools==1.7.0 onnx==1.17.0

这些包的作用分别是：

lightgbm: 用于训练梯度提升树模型
onnxmltools: 将模型转换为 ONNX 格式
onnx: ONNX 运行时支持

数据加载与预处理

我们使用一个企业破产预测的示例数据集：

df = (
    spark.read.format("csv")
    .option("header", True)
    .option("inferSchema", True)
    .load(
        "wasbs://publicwasb@mmlspark.blob.core.windows.net/company_bankruptcy_prediction_data.csv"
    )
)

display(df)

数据集包含多个财务指标，如：

Interest Coverage Ratio (利息保障倍数)
Net Income Flag (净利润标志)
Equity to Liability (权益负债比)

训练 LightGBM 模型

使用 SynapseML 提供的 LightGBM 分类器进行模型训练：

from pyspark.ml.feature import VectorAssembler
from synapse.ml.lightgbm import LightGBMClassifier

# 特征工程
feature_cols = df.columns[1:]
featurizer = VectorAssembler(inputCols=feature_cols, outputCol="features")
train_data = featurizer.transform(df)["Bankrupt?", "features"]

# 模型配置
model = (
    LightGBMClassifier(featuresCol="features", labelCol="Bankrupt?")
    .setEarlyStoppingRound(300)
    .setLambdaL1(0.5)
    .setNumIterations(1000)
    .setNumThreads(-1)
    .setMaxDeltaStep(0.5)
    .setNumLeaves(31)
    .setMaxDepth(-1)
    .setBaggingFraction(0.7)
    .setFeatureFraction(0.7)
    .setBaggingFreq(2)
    .setObjective("binary")
    .setIsUnbalance(True)
    .setMinSumHessianInLeaf(20)
    .setMinGainToSplit(0.01)
)

# 模型训练
model = model.fit(train_data)

模型转换：从 LightGBM 到 ONNX

ONNX (Open Neural Network Exchange) 是一种开放的模型表示格式，支持跨框架的模型互操作。我们将训练好的 LightGBM 模型转换为 ONNX 格式：

import lightgbm as lgb
from lightgbm import Booster, LGBMClassifier

def convertModel(lgbm_model: LGBMClassifier or Booster, input_size: int) -> bytes:
    from onnxmltools.convert import convert_lightgbm
    from onnxconverter_common.data_types import FloatTensorType

    initial_types = [("input", FloatTensorType([-1, input_size]))]
    onnx_model = convert_lightgbm(
        lgbm_model, initial_types=initial_types, target_opset=9
    )
    return onnx_model.SerializeToString()

# 获取模型并转换
booster_model_str = model.getLightGBMBooster().modelStr().get()
booster = lgb.Booster(model_str=booster_model_str)
model_payload_ml = convertModel(booster, len(feature_cols))

使用 ONNXModel 进行分布式推理

SynapseML 提供了 ONNXModel 类来方便地在 Spark 上运行 ONNX 模型：

from synapse.ml.onnx import ONNXModel

# 初始化ONNX模型
onnx_ml = ONNXModel().setModelPayload(model_payload_ml)

# 配置输入输出映射
onnx_ml = (
    onnx_ml.setDeviceType("CPU")
    .setFeedDict({"input": "features"})
    .setFetchDict({"probability": "probabilities", "prediction": "label"})
    .setMiniBatchSize(5000)
)

生成测试数据并推理

为了演示大规模推理能力，我们生成一个包含100万条记录的测试数据集：

from pyspark.ml.feature import VectorAssembler
import pandas as pd
import numpy as np

n = 1000 * 1000  # 100万条记录
m = 95           # 特征数量
test = np.random.rand(n, m)
testPdf = pd.DataFrame(test)
cols = list(map(str, testPdf.columns))
testDf = spark.createDataFrame(testPdf)
testDf = testDf.union(testDf).repartition(200)  # 扩大数据量并重新分区
testDf = (
    VectorAssembler()
    .setInputCols(cols)
    .setOutputCol("features")
    .transform(testDf)
    .drop(*cols)
    .cache()
)

# 执行推理
display(onnx_ml.transform(testDf))