Partilhar via

Executar inferência em lote usando um Spark DataFrame

Este artigo descreve como executar inferência em lote em um Spark DataFrame usando um modelo registrado no Databricks. O fluxo de trabalho se aplica a vários modelos de aprendizado de máquina e aprendizado profundo, incluindo TensorFlow, PyTorch e scikit-learn. Inclui práticas recomendadas para carregamento de dados, inferência de modelo e ajuste de desempenho.

Para inferência de modelo para aplicativos de aprendizado profundo, o Azure Databricks recomenda o seguinte fluxo de trabalho. Por exemplo, notebooks que usam TensorFlow e PyTorch, consulte Exemplos de inferência em lote.

Fluxo de trabalho de inferência de modelo

O Databricks recomenda o seguinte fluxo de trabalho para executar inferência em lote usando o Spark DataFrames.

Etapa 1: Configuração do ambiente

Certifique-se de que seu cluster execute uma versão compatível do Databricks ML Runtime para corresponder ao ambiente de treinamento. O modelo, registrado usando MLflow, contém os requisitos que podem ser instalados para garantir que os ambientes de treinamento e inferência coincidam.

requirements_path = os.path.join(local_path, "requirements.txt")
if not os.path.exists(requirements_path):
    dbutils.fs.put("file:" + requirements_path, "", True)

%pip install -r $requirements_path
%restart_python

Etapa 2: Carregar dados no Spark DataFrames

Dependendo do tipo de dados, use o método apropriado para carregar dados em um Spark DataFrame:

Tipo de dados Método
Tabela do Catálogo Unity (Recomendado) table = spark.table(input_table_name)
Arquivos de imagem (JPG, PNG) files_df = spark.createDataFrame(map(lambda path: (path,), file_paths), ["path"])
TFRecords df = spark.read.format("tfrecords").load(image_path)
Outros formatos (Parquet, CSV, JSON, JDBC) Carregue usando fontes de dados do Spark.

Etapa 3: Carregar o modelo do registo de modelos

Este exemplo usa um modelo do Databricks Model Registry para inferência.

predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri)

Etapa 4: Executar a inferência do modelo usando Pandas UDFs

As UDFs de Pandas utilizam o Apache Arrow para transferência eficiente de dados e o pandas para processamento. As etapas típicas da inferência com pandas UDFs são:

  1. Carregue o modelo treinado: utilize o MLflow para criar um Spark UDF para realizar inferência.
  2. Pré-processar dados de entrada: verifique se o esquema de entrada corresponde aos requisitos do modelo.
  3. Executar previsão de modelo: use a função UDF do modelo no DataFrame.
df_result = df_spark.withColumn("prediction", predict_udf(*df_spark.columns))
  1. (Recomendado) Salve previsões no Catálogo Unity.

O exemplo a seguir salva previsões no Unity Catalog.

df_result.write.mode("overwrite").saveAsTable(output_table)

Ajuste de desempenho para inferência de modelo

Esta seção fornece algumas dicas para depuração e ajuste de desempenho para inferência de modelo no Azure Databricks. Para obter uma visão geral, consulte realizar inferência em lote usando um DataFrame do Spark.

Normalmente, há duas partes principais na inferência do modelo: pipeline de entrada de dados e inferência do modelo. O pipeline de entrada de dados é exigente em entrada/saída de dados e a inferência do modelo é intensiva em computação. Determinar o gargalo do fluxo de trabalho é simples. Aqui estão algumas abordagens:

  • Reduza o modelo a um modelo trivial e meça os exemplos por segundo. Se a diferença do tempo de ponta a ponta entre o modelo completo e o modelo trivial for mínima, então o pipeline de entrada de dados é provavelmente um gargalo, caso contrário, a inferência do modelo é o gargalo.
  • Se estiver a executar a inferência de modelo com GPU, verifique as métricas de utilização da GPU . Se a utilização da GPU não for consistentemente alta, isso pode ser uma indicação de que o pipeline de entrada de dados é o gargalo.

Otimize o pipeline de entrada de dados

O uso de GPUs pode otimizar eficientemente a velocidade de execução para inferência de modelos. À medida que GPUs e outros aceleradores se tornam mais rápidos, é importante que o pipeline de entrada de dados acompanhe a demanda. O pipeline de entrada de dados lê os dados nos Spark DataFrames, transforma-os e carrega-os como a entrada para a inferência do modelo. Se a entrada de dados for o gargalo, aqui estão algumas dicas para aumentar a taxa de transferência de E/S:

  • Defina os registros máximos por lote. Um número maior de registos máximos pode reduzir a sobrecarga de E/S para chamar a função UDF, desde que os registos caibam na memória. Para definir o tamanho do lote, defina a seguinte configuração:

    spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "5000")

  • Carregue os dados em lotes e pré-busque-os ao pré-processar os dados de entrada no UDF pandas.

    Para o TensorFlow, o Azure Databricks recomenda o uso da API tf.data. Você pode analisar o mapa em paralelo definindo num_parallel_calls numa função map e chamar prefetch e batch para pré-carregamento e processamento em lotes.

    dataset.map(parse_example, num_parallel_calls=num_process).prefetch(prefetch_size).batch(batch_size)

    Para o PyTorch, o Azure Databricks recomenda o uso da classe DataLoader. Você pode definir batch_size para envio em lote e num_workers para carregamento de dados paralelos.

    torch.utils.data.DataLoader(images, batch_size=batch_size, num_workers=num_process)

Exemplos de inferência em lote

Os exemplos nesta seção seguem o fluxo de trabalho de inferência de aprendizado profundo recomendado. Estes exemplos ilustram como executar a inferência do modelo usando um modelo de rede neural pré-treinado de redes residuais profundas (ResNets).

Extração de dados estruturados e inferência em lote usando o Spark UDF

O bloco de anotações de exemplo a seguir demonstra o desenvolvimento, o registro em log e a avaliação de um agente simples para extração de dados estruturados para transformar dados brutos e não estruturados em informações organizadas e utilizáveis por meio de técnicas de extração automatizada. Essa abordagem demonstra como implementar agentes personalizados para inferência em lote usando a classe PythonModel do MLflow e empregar o modelo de agente registrado como uma função de User-Defined de faísca (UDF). Este caderno também mostra como aproveitar o Mosaic AI Agent Evaluation para avaliar a precisão usando dados de verdade do solo.

Extração de dados estruturados e inferência em lote usando o Spark UDF

Obter caderno