ONNX-inferens på Spark

ONNX (Open Neural Network Exchange) leverer en portabel, hardware-optimeret runtime til maskinlæringsmodeller. Ved at konvertere en model til ONNX-format kan du køre batch-inferens på Spark med lavere latenstid og uden at være afhængig af det oprindelige træningsrammeværk ved forudsigelsestidspunktet.

I denne artikel træner du en LightGBM-model med SynapseML, konverterer den til ONNX-format og bruger derefter ONNX-modellen til at udføre inferenser på Spark i Microsoft Fabric.

Forudsætninger

• Få et Microsoft Fabric-abonnement. Du kan også tilmelde dig en gratis Prøveversion af Microsoft Fabric.

• Log på Microsoft Fabric.

• Skift til Fabric ved at bruge experience-switcheren nederst til venstre på din startside.

  

• Vedhæft din notesbog til et lakehouse. På venstre side af din notesbog vælger du Tilføj for at tilføje et eksisterende søhus eller oprette et.
• Fabric Runtime 1.2 eller senere.

Installer påkrævede pakker

Kør følgende felt i din bærbare for at installere de nødvendige pakker. onnxmltools-pakken er ikke forudinstalleret i Fabric-runtime.

%pip install onnxmltools --quiet

Når installationen er færdig, skal du sikre dig, at pakkerne er tilgængelige:

import onnxmltools
import lightgbm
print(f"onnxmltools version: {onnxmltools.__version__}")
print(f"lightgbm version: {lightgbm.__version__}")

Indlæs eksempeldataene

Load the bankruptcy prediction dataset from public Azure Blob Storage:

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

print(f"Rows: {df.count()}, Columns: {len(df.columns)}")
display(df.limit(5))

Den viste tabel indeholder kolonner såsom:

Træn en LightGBM-model

Brug kolonnerne VectorAssembler til at kombinere feature, og træn derefter en LightGBMClassifier:

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?")
    .setDataTransferMode("bulk")
    .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)

Verificér den model, der er trænet med succes:

print(f"Model type: {type(model).__name__}")
print(f"Number of features: {len(feature_cols)}")

Konvertér modellen til ONNX-format

Eksporter den trænede model til en LightGBM-booster, og konverter den derefter til ONNX:

import lightgbm as lgb
from typing import Union
from lightgbm import Booster, LGBMClassifier
from onnxmltools.convert import convert_lightgbm
from onnxmltools.convert.common.data_types import FloatTensorType


def convert_to_onnx(lgbm_model: Union[LGBMClassifier, Booster], input_size: int) -> bytes:
    initial_types = [("input", FloatTensorType([-1, input_size]))]
    onnx_model = convert_lightgbm(
        lgbm_model, initial_types=initial_types, target_opset=13
    )
    return onnx_model.SerializeToString()


booster_model_str = model.getLightGBMBooster().modelStr().get()
booster = lgb.Booster(model_str=booster_model_str)
model_payload_ml = convert_to_onnx(booster, len(feature_cols))

Kontroller at ONNX-konverteringen lykkedes:

print(f"ONNX model payload size: {len(model_payload_ml)} bytes")
assert len(model_payload_ml) > 0, "ONNX conversion failed: empty payload"

Outputtet viser ONNX-modellens payload-størrelse i bytes (typisk omkring 800.000 bytes).

Indlæs og konfigurer ONNX-modellen

Indlæs ONNX-payloaden i en SynapseML ONNXModel og inspicer modellens input og output:

from synapse.ml.onnx import ONNXModel

onnx_ml = ONNXModel().setModelPayload(model_payload_ml)

print("Model inputs:" + str(onnx_ml.getModelInputs()))
print("Model outputs:" + str(onnx_ml.getModelOutputs()))

Outputtet viser modellens input- og outputnoder.

Konfigurér modellen ved at mappe input- og outputkolonner. Den FeedDict kortlægger ONNX-modellens inputnavne til DataFrame-kolonnenavne. Kortlægningen FetchDict af ønskede outputkolonnenavne til ONNX-modellens outputnavne:

onnx_ml = (
    onnx_ml.setDeviceType("CPU")
    .setFeedDict({"input": "features"})
    .setFetchDict({"probability": "probabilities", "prediction": "label"})
    .setMiniBatchSize(5000)
)

Løbinferens

Opret testdata og transformér dem gennem ONNX-modellen:

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

n = 10000
m = 95
test = np.random.rand(n, m)
testPdf = pd.DataFrame(test)
cols = list(map(str, testPdf.columns))
testDf = spark.createDataFrame(testPdf)
testDf = testDf.repartition(4)
testDf = (
    VectorAssembler()
    .setInputCols(cols)
    .setOutputCol("features")
    .transform(testDf)
    .drop(*cols)
    .cache()
)

display(onnx_ml.transform(testDf))

Outputtet skal indeholde kolonner for features, prediction, og probability:

Verificér den resulterende slutning:

results = onnx_ml.transform(testDf)
print(f"Result count: {results.count()}")
print(f"Output columns: {results.columns}")
assert "prediction" in results.columns, "Missing prediction column"
assert "probability" in results.columns, "Missing probability column"

Outputtet bekræfter, at alle testrækker blev scoret, og at resultatdatarammen indeholder features, prediction, og probability kolonnerne.

Fejlfinding

Ryd op i ressourcer

Hvis du ikke længere har brug for den cachede test-DataFrame, så genoplever den for at frigøre klyngehukommelsen:

testDf.unpersist()

Relateret indhold

• Sådan bruges Kerne-SHAP til at forklare en tabelklassificeringsmodel
• Sådan bruger du SynapseML til registrering af flerdimensionel uregelmæssighed
• Sådan bygger du en søgemaskine med SynapseML