ONNX-slutsatsdragning på Spark

ONNX (Open Neural Network Exchange) tillhandahåller en portabel, maskinvaruoptimerad körning för maskininlärningsmodeller. Genom att konvertera en modell till ONNX-format kan du köra batchinferens i Spark med lägre latens och utan att vara beroende av det ursprungliga träningsramverket vid prediktionstillfället.

I den här artikeln tränar du en LightGBM-modell med SynapseML, konverterar den till ONNX-format och använder sedan ONNX-modellen för att utföra slutsatsdragning på Spark i Microsoft Fabric.

Förutsättningar

  • Bifoga anteckningsboken till ett sjöhus. Till vänster i anteckningsboken väljer du Lägg till för att lägga till ett befintligt sjöhus eller skapa ett.
  • Fabric Runtime 1.2 eller senare.

Installera de paket som krävs

Kör följande cell i din notebook för att installera de nödvändiga paketen. paketet onnxmltools är inte förinstallerat i Fabric-körningen.

%pip install onnxmltools --quiet

När installationen är klar kontrollerar du att paketen är tillgängliga:

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

Note

paketet lightgbm är förinstallerat i Fabric Runtime 1.2 och senare. Du behöver bara installera onnxmltools.

Läs in exempeldata

Ladda in datauppsättningen för konkursprognoser från offentlig 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 visade tabellen innehåller kolumner som:

Konkurs? Nettoinkomstindikator Förhållande mellan eget kapital och skuld
0 1.0 0.0165
0 1.0 0.0208

Träna en LightGBM-modell

Använd VectorAssembler för att kombinera funktionskolumner och träna sedan 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)

Kontrollera att modellen har tränats framgångsrikt:

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

Konvertera modellen till ONNX-format

Exportera den tränade modellen till en LightGBM-booster och konvertera den sedan till 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))

Kontrollera att ONNX-konverteringen lyckades:

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

Utdata visar ONNX-modellens nyttolaststorlek i byte (vanligtvis cirka 800 000 byte).

Important

Används from onnxmltools.convert.common.data_types import FloatTensorType för typdefinitionen. Den äldre importsökvägen from onnxconverter_common.data_types import FloatTensorType är inte kompatibel med aktuella versioner av onnxmltools.

Läsa in och konfigurera ONNX-modellen

Läs in ONNX-nyttolasten i en SynapseML ONNXModel och inspektera modellens indata och utdata:

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()))

Utdata visar modellens indata- och utdatanoder.

Konfigurera modellen genom att mappa indata- och utdatakolumner. FeedDict mappar ONNX-modellens indatanamn till DataFrame-kolumnnamn. FetchDict mappar önskade namn på utdatakolumner till utdatanamn i ONNX-modellen:

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

Kör slutsatsdragning

Skapa testdata och transformera dem via 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))

Note

Eftersom testdata genereras slumpmässigt representerar förutsägelsevärdena inte verkliga resultat. Det här avsnittet visar att ONNX-modellen körs korrekt på Spark.

Utdata ska innehålla kolumner för features, predictionoch probability:

Egenskaper prognos sannolikhet
{"type":1,"values":[0.105... 0 {"0":0.835...
{"type":1,"values":[0.814... 0 {"0":0.658...

Kontrollera resultatet av slutsatsdragningen:

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"

Utdata bekräftar att alla testrader har poängsatts och resultatet DataFrame innehåller kolumnerna features, predictionoch probability .

Felsökning

Problematik Orsak Lösning
ModuleNotFoundError: No module named 'onnxmltools' Paketet är inte förinstallerat i Fabric-körmiljön. Kör %pip install onnxmltools --quiet och starta om Python kernel.
RuntimeError: Operator LgbmClassifier got an input with a wrong type Fel importsökväg för FloatTensorType. Använd from onnxmltools.convert.common.data_types import FloatTensorType i stället för att importera från onnxconverter_common.data_types.
ModuleNotFoundError: No module named 'onnx.mapping' Inkompatibel onnxmltools version 1.7.0 eller tidigare med det aktuella onnx paketet. Kör %pip install onnxmltools --upgrade --quiet för att installera en kompatibel version.
ONNX conversion returns empty payload Extrahering av boostermodellsträng misslyckades. Kontrollera att model.getLightGBMBooster().modelStr().get() returnerar en icke-tom sträng före konvertering.
Feature (Column_) appears more than one time Under model.fit() Kolumner i datauppsättningar med specialtecken ger upphov till dubblettnamn efter LightGBM:s sanering. Lägg till .setDataTransferMode("bulk") i konfigurationen LightGBMClassifier . Massläge använder Apache Arrow och undviker problemet med kolumnnamnssanering.
AssertionError för SparkContext i ONNXModel() Spark-sessionen initieras inte. Kör den här koden i en Fabric-notebook med ett anslutet lakehouse. Variabeln spark förinitieras av körmiljön.

Rensa resurser

Om du inte längre behöver test-DataFramen som har cachats, tar du bort den från minnet för att frigöra klustrets minne:

testDf.unpersist()