Inférence ONNX sur Spark

ONNX (Open Neural Network Exchange) fournit un runtime portable optimisé pour le matériel pour les modèles Machine Learning. En convertissant un modèle au format ONNX, vous pouvez exécuter l’inférence par lots sur Spark avec une latence inférieure et sans dépendre de l’infrastructure d’apprentissage d’origine au moment de la prédiction.

Dans cet article, vous allez entraîner un modèle LightGBM avec SynapseML, le convertir au format ONNX, puis utiliser le modèle ONNX pour effectuer une inférence sur Spark dans Microsoft Fabric.

Conditions préalables

• Obtenez un abonnement Microsoft Fabric. Ou, inscrivez-vous pour un essai gratuit de Microsoft Fabric.

• Connectez-vous à Microsoft Fabric.

• Basculez vers Fabric à l’aide du sélecteur d’expérience situé en bas à gauche de votre page d’accueil.

  

• Fixez votre notebook à un lakehouse. Sur le côté gauche de votre bloc-notes, sélectionnez Ajouter pour ajouter un lakehouse existant ou en créer un.
• Fabric Runtime 1.2 ou version ultérieure.

Installer les packages requis

Exécutez la cellule suivante dans votre notebook pour installer les packages requis. Le package onnxmltools n'est pas préinstallé dans le runtime Fabric.

%pip install onnxmltools --quiet

Une fois l’installation terminée, vérifiez que les packages sont disponibles :

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

Charger les exemples de données

Chargez le jeu de données de prédiction de faillite à partir du Stockage Blob Azure public :

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

Le tableau affiché comprend des colonnes telles que :

Entraîner un modèle LightGBM

Utilisez la fonction VectorAssembler pour combiner des colonnes de caractéristiques, puis entraînez un 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)

Vérifiez que le modèle a été correctement formé :

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

Convertir le modèle au format ONNX

Exportez le modèle entraîné vers un booster LightGBM, puis convertissez-le en 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))

Vérifiez que la conversion ONNX a réussi :

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

La sortie affiche la taille de charge utile du modèle ONNX en octets (généralement environ 800 000 octets).

Charger et configurer le modèle ONNX

Chargez la charge utile ONNX dans un SynapseML ONNXModel et inspectez les entrées et sorties du modèle :

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

La sortie répertorie les nœuds d’entrée et de sortie du modèle.

Configurez le modèle en mappant les colonnes d’entrée et de sortie. Le FeedDict associe les noms d’entrée du modèle ONNX aux noms de colonnes du DataFrame. Le FetchDict associe les noms de colonnes de sortie souhaités aux noms de sortie du modèle ONNX :

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

Exécuter l’inférence

Créez des données de test et transformez-les via le modèle ONNX :

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

La sortie doit contenir des colonnes pour features, predictionet probability:

Vérifiez les résultats générés par l’inférence :

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"

La sortie confirme que toutes les lignes de test ont été évaluées et que le DataFrame de résultat contient les colonnes features, prediction et probability.

Résolution des problèmes

Nettoyer les ressources

Si vous n’avez plus besoin du DataFrame de test mis en cache, annulez-le pour libérer de la mémoire du cluster :

testDf.unpersist()

Contenu connexe

• Comment utiliser Kernel SHAP pour expliquer un modèle de classification tabulaire
• Comment utiliser SynapseML pour la détection d’anomalies multivariées
• Comment créer un moteur de recherche avec SynapseML