Inferência ONNX no Spark

ONNX (Open Neural Network Exchange) fornece um tempo de execução portátil e otimizado para hardware para modelos de aprendizagem automática. Ao converter um modelo para formato ONNX, pode executar inferência em lote no Spark com menor latência e sem depender do framework de treino original no momento da previsão.

Neste artigo, treina um modelo LightGBM com SynapseML, converte-o para o formato ONNX e depois usa o modelo ONNX para realizar inferência no Spark no Microsoft Fabric.

Pré-requisitos

• Obtenha uma assinatura do Microsoft Fabric. Ou inscreva-se para uma avaliação gratuita do Microsoft Fabric.

• Faça login no Microsoft Fabric.

• Mude para o Fabric usando o seletor de experiência no canto inferior esquerdo da sua página inicial.

  

• Ligue o seu bloco de notas a uma casa no lago. No lado esquerdo do seu caderno, selecione Adicionar para adicionar uma casa de lago existente ou criar uma.
• Fabric Runtime 1.2 ou posterior.

Instalar pacotes necessários

Execute a célula seguinte no seu portátil para instalar os pacotes necessários. O pacote onnxmltools não está pré-instalado no tempo de execução Fabric.

%pip install onnxmltools --quiet

Após a instalação concluída, verifique se os pacotes estão disponíveis:

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

Carregue os dados de exemplo

Carregue o conjunto de dados de previsão de falências a partir do Armazenamento de Blobs do Azure 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))

A tabela apresentada inclui colunas como:

Treine um modelo LightGBM

Use o VectorAssembler para combinar colunas de características, depois treine um 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)

Verifique se o modelo foi treinado com sucesso:

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

Converter o modelo para o formato ONNX

Exporta o modelo treinado para um booster LightGBM e depois converte-o para 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))

Verifique se a conversão ONNX foi bem-sucedida:

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

A saída mostra o tamanho da carga útil do modelo ONNX em bytes (tipicamente cerca de 800.000 bytes).

Carregar e configurar o modelo ONNX

Carregue a carga útil ONNX num SynapseML ONNXModel e inspecione as entradas e saídas do 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()))

A saída lista os nós de entrada e saída do modelo.

Configure o modelo mapeando as colunas de entrada e saída. O FeedDict associa os nomes de entrada do modelo ONNX aos nomes de colunas do DataFrame. O FetchDict associa os nomes das colunas de saída pretendidas aos nomes de saída do modelo ONNX:

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

Executar inferência

Crie dados de teste e transforme-os através do 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))

A saída deve conter colunas para features, prediction, e probability:

Verifique os resultados da inferência produzida:

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"

A saída confirma que todas as linhas de teste foram classificadas e o DataFrame de resultados contém as colunas features, prediction e probability.

Troubleshooting

Limpeza de recursos

Se já não precisar do DataFrame de teste em cache, remova-o da cache para libertar memória do cluster:

testDf.unpersist()

Conteúdo relacionado

• Como usar o Kernel SHAP para explicar um modelo de classificação tabular
• Como usar SynapseML para deteção de anomalias multivariadas
• Como construir um motor de busca com o SynapseML