Odvození ONNX ve Sparku

  • Článek

V tomto příkladu vytrénujete model LightGBM a převedete ho na formát ONNX . Po převodu použijete model k odvození některých testovacích dat ve Sparku.

V tomto příkladu se používají následující balíčky a verze Pythonu:

  • onnxmltools==1.7.0
  • lightgbm==3.2.1

Požadavky

  • Připojte poznámkový blok k jezeru. Na levé straně vyberte Přidat a přidejte existující jezerní dům nebo vytvořte jezero.
  • Možná budete muset nainstalovat onnxmltools přidáním !pip install onnxmltools==1.7.0 do buňky kódu a následným spuštěním buňky.

Načtení ukázkových dat

Pokud chcete načíst ukázková data, přidejte do buněk v poznámkovém bloku následující příklady kódu a pak tyto buňky spusťte:

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)

Výstup by měl vypadat podobně jako v následující tabulce, i když se hodnoty a počet řádků můžou lišit:

Poměr pokrytí úroku Příznak Čistý příjem Vlastní kapitál vůči odpovědnosti
0.5641 1.0 0.0165
0.5702 1.0 0.0208
0.5673 1.0 0.0165

Trénování modelu pomocí 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)

Převod modelu do formátu ONNX

Následující kód exportuje trénovaný model do boosteru LightGBM a pak ho převede do formátu 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))

Po převodu načtěte datovou část ONNXModel ONNX a zkontrolujte vstupy a výstupy modelu:

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

Namapujte vstup modelu na název sloupce vstupního datového rámce (FeedDict) a namapujte názvy sloupců výstupního datového rámce na výstupy modelu (FetchDict).

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

Použití modelu k odvozování

Pokud chcete provést odvozování s modelem, následující kód vytvoří testovací data a transformuje data prostřednictvím modelu ONNX.

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

Výstup by měl vypadat podobně jako v následující tabulce, i když se hodnoty a počet řádků můžou lišit:

Index Funkce Predikce Probability
0 "{"type":1,"values":[0.105... 0 "{"0":0.835...
2 "{"type":1,"values":[0.814... 0 "{"0":0.658...

Související obsah