Inferencia de ONNX en Spark

ONNX (Open Neural Network Exchange) proporciona un entorno de ejecución portátil optimizado para hardware para modelos de aprendizaje automático. Al convertir un modelo en formato ONNX, puede ejecutar la inferencia por lotes en Spark con menor latencia y sin depender del marco de entrenamiento original en el momento de la predicción.

En este artículo, entrena un modelo LightGBM con SynapseML, lo convierte al formato ONNX y, a continuación, usa el modelo ONNX para realizar inferencias en Spark en Microsoft Fabric.

Prerrequisitos

• Obtenga una suscripción a Microsoft Fabric. O bien, regístrese para obtener una evaluación gratuita de Microsoft Fabric.

• Inicie sesión en Microsoft Fabric.

• Cambie a Fabric mediante el conmutador de experiencia en el lado inferior izquierdo de la página principal.

  

• Adjunte su cuaderno a un almacén de lago de datos. En el lado izquierdo del cuaderno, seleccione Agregar para agregar un lago existente o crear uno.
• Fabric Runtime 1.2 o posterior.

Instalación de paquetes necesarios

Ejecute la siguiente celda en el cuaderno para instalar los paquetes necesarios. El paquete onnxmltools no está preinstalado en el entorno de ejecución de Fabric.

%pip install onnxmltools --quiet

Una vez completada la instalación, compruebe que los paquetes están disponibles:

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

Carga de los datos de ejemplo

Cargue el conjunto de datos de predicción de quiebras del Azure Blob Storage público:

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

La tabla mostrada incluye columnas como:

Entrena un modelo LightGBM

Use el VectorAssembler para combinar columnas de características y, a continuación, entrene 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)

Compruebe que el modelo se entrenó correctamente:

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

Convertir el modelo en formato ONNX

Exporte el modelo entrenado a un booster de LightGBM y, a continuación, conviértalo a 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))

Compruebe que la conversión de ONNX se realizó correctamente:

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

La salida muestra el tamaño de carga del modelo ONNX en bytes (normalmente alrededor de 800 000 bytes).

Carga y configuración del modelo ONNX

Cargue la carga de ONNX en synapseML ONNXModel e inspeccione las entradas y salidas del modelo:

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 salida enumera los nodos de entrada y salida del modelo.

Configure el modelo mediante la asignación de columnas de entrada y salida. El FeedDict asigna los nombres de entrada del modelo ONNX a las columnas del DataFrame. El FetchDict asigna los nombres deseados de las columnas de salida a los nombres de salida del modelo ONNX:

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

Ejecución de la inferencia

Crear datos de prueba y transformarlos a través del modelo 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 salida debe contener columnas para features, predictiony probability:

Compruebe los resultados generados por la inferencia:

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 salida confirma que todas las filas de prueba se puntuaron y el resultado DataFrame contiene las featurescolumnas , predictiony probability .

Solución de problemas

Limpieza de recursos

Si ya no necesita el DataFrame de prueba almacenado en caché, despersistelo para liberar memoria del clúster:

testDf.unpersist()

Contenido relacionado

• Cómo usar Kernel SHAP para explicar un modelo de clasificación tabular
• Uso de SynapseML para la detección de anomalías multivariante
• Creación de un motor de búsqueda con SynapseML