Registrar modelos de MLflow

  • Artigo

Este artigo descreve como registrar seus modelos treinados (ou artefatos) como modelos do MLflow. Ele explora as diferentes maneiras de personalizar como o MLflow empacota seus modelos e como ele executa esses modelos.

Por que registrar modelos em vez de artefatos?

De artefatos a modelos no MLflow descreve a diferença entre o registro em log de artefatos ou arquivos, em comparação com o registro em log de modelos do MLflow.

Um modelo do MLflow também é um artefato. No entanto, esse modelo tem uma estrutura específica que serve como um contrato entre a pessoa que criou o modelo e a pessoa que pretende usá-lo. Esse contrato ajuda a construir uma ponte entre os próprios artefatos e seus significados.

O registro em log de modelos tem estas vantagens:

  • Você pode carregar modelos diretamente, para inferência, com mlflow.<flavor>.load_model, e você pode usar a função predict
  • As entradas de pipeline podem usar modelos diretamente
  • Você pode implantar modelos sem indicação de um script de pontuação ou de um ambiente
  • O Swagger é habilitado automaticamente em pontos de extremidade implantados e o Estúdio do Azure Machine Learning pode usar o recurso Teste
  • Você pode usar o painel de IA Responsável

Esta seção descreve como usar o conceito do modelo no Azure Machine Learning com o MLflow:

Registrar modelos usando o log automático

Você pode usar a funcionalidade de registro em log automático do MLflow. O registro em log automático permite que o MLflow instrua a estrutura em uso para registrar em log todas as métricas, parâmetros, artefatos e modelos que a estrutura considere relevantes. Por padrão, se o registro em log automático estiver habilitado, a maioria dos modelos será registrado em log. Em algumas situações, algumas variantes podem não registrar um modelo em log. Por exemplo, a variante PySpark não registra modelos que excedem um determinado tamanho.

Use mlflow.autolog() ou mlflow.<flavor>.autolog() para ativar o registro em log automático. Este exemplo usa autolog() para registrar um modelo de classificador treinado com o XGBoost:

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score

mlflow.autolog()

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)

y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)

Dica

Se usar pipelines do Machine Learning, por exemplo, pipelines do Scikit-Learn, use a funcionalidade autolog dessa variante de pipeline para registrar modelos. O registro em log de modelo ocorre automaticamente quando o método fit() é chamado no objeto de pipeline. Treinamento e o acompanhamento de um classificador XGBoost com o notebook do MLflow demonstra como registrar um modelo com pré-processamento usando pipelines.

Registrar modelos com uma assinatura, ambiente ou exemplos personalizados

O método mlflow.<flavor>.log_model do MLflow pode registrar manualmente modelos. Este fluxo de trabalho pode controlar diferentes aspectos do registro em log do modelo.

Use esse método quando:

  • Você deseja indicar pacotes pip ou um ambiente do Conda diferentes daqueles que são detectados automaticamente
  • Você deseja incluir exemplos de entrada
  • Você deseja incluir artefatos específicos no pacote necessário
  • autolog não infere corretamente a sua assinatura. Isso importa quando você lida com entradas de tensor, em que a assinatura precisa de formas específicas
  • O comportamento do registro em log automático não abrange o seu objetivo por algum motivo

Este exemplo de código registra um modelo para um classificador XGBoost:

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature
from mlflow.utils.environment import _mlflow_conda_env

mlflow.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_pred = model.predict(X_test)

accuracy = accuracy_score(y_test, y_pred)

# Signature
signature = infer_signature(X_test, y_test)

# Conda environment
custom_env =_mlflow_conda_env(
    additional_conda_deps=None,
    additional_pip_deps=["xgboost==1.5.2"],
    additional_conda_channels=None,
)

# Sample
input_example = X_train.sample(n=1)

# Log the model manually
mlflow.xgboost.log_model(model, 
                         artifact_path="classifier", 
                         conda_env=custom_env,
                         signature=signature,
                         input_example=input_example)

Observação

  • autolog tem a configuração log_models=False. Isso impede o registro automático do modelo do MLflow. O registro automático do modelo do MLflow ocorre mais tarde, como um processo manual
  • Use o método infer_signature para tentar inferir a assinatura diretamente de entradas e saídas
  • O método mlflow.utils.environment._mlflow_conda_env é um método privado no SDK do MLflow. Neste exemplo, ele torna o código mais simples, mas use-o com cuidado. Isso pode mudar no futuro. Como alternativa, você pode gerar a definição YAML manualmente como um dicionário do Python.

Registro de modelos com um comportamento diferente no método de previsão

Ao registrar em log um modelo com mlflow.autolog ou mlflow.<flavor>.log_model, a variante do modelo determina como executar a inferência e o que o modelo retorna. O MLflow não impõe nenhum comportamento específico sobre a geração de resultados predict. Em alguns cenários, talvez você queira fazer algum pré-processamento ou pós-processamento antes e depois da execução do modelo.

Nesta situação, implemente pipelines de aprendizado de máquina que se movam diretamente de entradas para saídas. Embora esta implementação seja possível e, às vezes, incentivada para melhorar o desempenho, ela pode se tornar um desafio a ser alcançado. Nesses casos, personalizar como o seu modelo lida com a inferência pode ajudar, conforme explicado na próxima seção.

Registro de modelos personalizados

O MLflow dá suporte a muitas estruturas de machine learning, incluindo

  • CatBoost
  • FastAI
  • h2o
  • Keras
  • LightGBM
  • MLeap
  • MXNet Gluon
  • ONNX
  • Prophet
  • PyTorch
  • Scikit-Learn
  • spaCy
  • Spark MLLib
  • statsmodels
  • TensorFlow
  • XGBoost

No entanto, talvez seja necessário alterar a maneira como uma variante funciona, registrar um modelo sem suporte nativo pelo MLflow ou até mesmo registrar um modelo que use vários elementos de estruturas diferentes. Nesses casos, talvez seja necessário criar uma variante de modelo personalizado.

Para resolver o problema, o MLflow apresenta a variante pyfunc (partindo de uma função Python). Essa variante pode registrar qualquer objeto como um modelo, desde que esse objeto atenda a duas condições:

  • Você implementa o método predict do método, pelo menos
  • O objeto Python herda de mlflow.pyfunc.PythonModel

Dica

Modelos serializáveis que implementam a API scikit-learn podem usar a variante Scikit-learn para registrar o modelo em log, independentemente de o modelo ter sido criado com o Scikit-learn. Se você puder manter o seu modelo no formato Pickle e o objeto tiver os métodos predict() e predict_proba() (pelo menos), você poderá usar mlflow.sklearn.log_model() para registrar o modelo em log dentro de uma execução do MLflow.

Se você criar um wrapper em torno do seu objeto de modelo existente, será mais simples de criar uma variante para o seu modelo personalizado. O MLflow o serializa e empacota para você. Os objetos do Python são serializáveis quando podem ser armazenados no sistema de arquivos como arquivos, geralmente no formato Pickle. Em runtime, o objeto pode se materializar a partir desse arquivo. Isso restaura todos os valores, propriedades e métodos disponíveis quando ele foi salvo.

Use esse método quando:

  • Você pode serializar o seu modelo no formato Pickle
  • Você deseja manter o estado do modelo, como ele era logo após o treinamento
  • Você deseja personalizar como a função predict funciona.

Este exemplo de código encapsula um modelo criado com o XGBoost, para que ele se comporte de forma diferente da implementação padrão da variante XGBoost. Em vez disso, ele retorna as probabilidades em vez das classes:

from mlflow.pyfunc import PythonModel, PythonModelContext

class ModelWrapper(PythonModel):
    def __init__(self, model):
        self._model = model

    def predict(self, context: PythonModelContext, data):
        # You don't have to keep the semantic meaning of `predict`. You can use here model.recommend(), model.forecast(), etc
        return self._model.predict_proba(data)

    # You can even add extra functions if you need to. Since the model is serialized,
    # all of them will be available when you load your model back.
    def predict_batch(self, data):
        pass

Registre um modelo personalizado na execução:

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature

mlflow.xgboost.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_probs = model.predict_proba(X_test)

accuracy = accuracy_score(y_test, y_probs.argmax(axis=1))
mlflow.log_metric("accuracy", accuracy)

signature = infer_signature(X_test, y_probs)
mlflow.pyfunc.log_model("classifier", 
                        python_model=ModelWrapper(model),
                        signature=signature)

Dica

Aqui, o método infer_signature usa y_probs para inferir a assinatura. Nossa coluna de destino tem a classe de destino, mas o modelo agora retorna as duas probabilidades para cada classe.

Próximas etapas