Dela via

ONNX-slutsatsdragning på Spark

  • Artikel

I det här exemplet tränar du en LightGBM-modell och konverterar modellen till ONNX-format . När du har konverterat använder du modellen för att härleda vissa testdata på Spark.

I det här exemplet används följande Python-paket och -versioner:

  • onnxmltools==1.7.0
  • lightgbm==3.2.1

Förutsättningar

  • Bifoga anteckningsboken till ett sjöhus. Till vänster väljer du Lägg till för att lägga till ett befintligt sjöhus eller skapa ett sjöhus.
  • Du kan behöva installera onnxmltools genom att lägga till !pip install onnxmltools==1.7.0 i en kodcell och sedan köra cellen.

Läs in exempeldata

Om du vill läsa in exempeldata lägger du till följande kodexempel i celler i anteckningsboken och kör sedan cellerna:

from pyspark.sql import SparkSession

# Bootstrap Spark Session
spark = SparkSession.builder.getOrCreate()

from synapse.ml.core.platform import *

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)

Utdata bör se ut ungefär som i följande tabell, men värdena och antalet rader kan skilja sig åt:

Räntetäckningsgrad Nettoinkomstflagga Eget kapital till skuld
0.5641 1.0 0.0165
0.5702 1.0 0.0208
0.5673 1.0 0.0165

Använda LightGBM för att träna en modell

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?", dataTransferMode="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)

Konvertera modellen till ONNX-format

Följande kod exporterar den tränade modellen till en LightGBM-booster och konverterar den sedan till ONNX-format:

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

Efter konverteringen läser du in ONNX-nyttolasten i en ONNXModel och inspekterar 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()))

Mappa modellindata till indataramens kolumnnamn (FeedDict) och mappa utdataramens kolumnnamn till modellutdata (FetchDict).

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

Använda modellen för slutsatsdragning

För att utföra slutsatsdragning med modellen skapar följande kod testdata och transformerar data via ONNX-modellen.

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

n = 1000 * 1000
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))

Utdata bör se ut ungefär som i följande tabell, men värdena och antalet rader kan skilja sig åt:

Index Funktioner Förutsägelse Sannolikhet
1 "{"type":1,"values":[0.105... 0 "{"0":0.835...
2 "{"type":1,"values":[0.814... 0 "{"0":0.658...

Relaterat innehåll