Registro de modelos de MLflow

2025-05-05

En este artículo se describe cómo registrar los modelos de aprendizaje automático entrenados o artefactos, como modelos de MLflow. MLflow es un marco de código abierto para administrar flujos de trabajo de aprendizaje automático. En este artículo se exploran varias opciones para personalizar la forma en que MLflow empaqueta y ejecuta modelos.

Prerrequisitos

El paquete del SDK mlflow de MLflow

¿Por qué registrar modelos en lugar de artefactos?

Un modelo de MLflow es un tipo de artefacto. Sin embargo, un modelo tiene una estructura específica que actúa como un contrato entre la persona que crea el modelo y la persona que pretende usarlo. Este contrato crea un puente entre los artefactos y sus significados.

Para entender la diferencia entre el registro de artefactos y archivos, y el registro de modelos de MLflow, consulte Artefactos y modelos en MLflow.

Puede registrar los archivos del modelo como artefactos, pero el registro de modelos ofrece las siguientes ventajas:

Puede usar mlflow.<flavor>.load_model para cargar directamente modelos para la inferencia y puede usar la predict función .
Las entradas de canalización pueden usar modelos directamente.
Puede implementar modelos sin especificar un script de puntuación o un entorno.
Swagger se activa automáticamente en los puntos finales desplegados. Como resultado, puede usar la característica de prueba en Azure Machine Learning Studio para probar modelos.
Puede usar el panel de Inteligencia Artificial Responsable. Para más información, consulte Uso del panel de inteligencia artificial responsable en Azure Machine Learning Studio.

Uso del registro automático para registrar modelos

Puede usar la funcionalidad de MLflow autolog para registrar automáticamente los modelos. Al usar el registro automático, MLflow captura todas las métricas, parámetros, artefactos y modelos pertinentes en el marco. Los datos registrados dependen del marco de trabajo. De forma predeterminada, si el registro automático está activado, la mayoría de los modelos se registran. En algunas situaciones, algunos tipos no registran modelos. Por ejemplo, el tipo de PySpark no registra modelos que superen un tamaño determinado.

Use mlflow.autolog o mlflow.<flavor>.autolog para activar el registro automático. El código siguiente usa autolog para registrar un modelo clasificador entrenado con 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)

Sugerencia

Si usa canalizaciones de aprendizaje automático, por ejemplo, canalizaciones scikit-learn, use la funcionalidad autolog de ese tipo de canalización para registrar modelos. El registro de modelos se ejecuta automáticamente al invocar el método fit en el objeto pipeline. Para un cuaderno que registra un modelo, incluye preprocesamiento y usa canalizaciones, consulte Entrenamiento y seguimiento de un clasificador XGBoost con MLflow.

Modelos de registro que usan una firma, un entorno o ejemplos personalizados

Puede usar el método MLflow mlflow.<flavor>.log_model para registrar manualmente los modelos. Este flujo de trabajo ofrece control sobre varios aspectos del registro de modelos.

Utilice este método cuando:

Quiere indicar un entorno de Conda o paquetes pip que difieren de los paquetes o entorno detectados automáticamente.
Quieres incluir ejemplos de entrada.
Quiere incluir artefactos específicos en el paquete que necesita.
El autolog método no deduce correctamente la firma. Este caso aparece cuando se trabaja con entradas de tensor, que requieren que la firma tenga una forma específica.
El autolog método no satisface todas sus necesidades.

El código siguiente registra un modelo de clasificador 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)

# Infer the signature.
signature = infer_signature(X_test, y_test)

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

# Sample the data.
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)

Nota:

La llamada a autolog usa una configuración de log_models=False. Esta configuración desactiva el registro automático de modelos de MLflow. El log_model método se usa más adelante para registrar manualmente el modelo.
El infer_signature método se usa para intentar deducir la firma directamente desde entradas y salidas.
mlflow.utils.environment._mlflow_conda_env es un método privado en el SDK de MLflow. En este ejemplo, simplifica el código. Pero use este método con precaución, ya que podría cambiar en el futuro. Como alternativa, puede generar la definición de YAML manualmente como diccionario de Python.

Modelos de registro que usan el comportamiento de predicción modificado

Cuando se usa mlflow.autolog o mlflow.<flavor>.log_model para registrar un modelo, el tipo de modelo determina cómo se realiza la inferencia. El tipo también determina lo que devuelve el modelo. MLflow no impone ningún comportamiento específico para la generación de resultados de predict. En algunos escenarios, es posible que desee preprocesar o posprocesar los datos.

En esta situación, puede implementar canalizaciones de aprendizaje automático que se mueven directamente de entradas a salidas. Aunque este tipo de implementación a veces puede mejorar el rendimiento, puede ser difícil de lograr. En tales casos, puede resultar útil personalizar cómo controla el modelo la inferencia. Para obtener más información, consulte la sección siguiente, Registro de modelos personalizados.

Registro de modelos personalizados

MLflow admite muchos marcos de aprendizaje automático, incluidos los siguientes tipos:

CatBoost
FastAI
H₂O
Keras
LightGBM
MLeap
ONNX
Profeta
PyTorch
scikit-learn
spaCy
Spark MLlib
statsmodels (paquete de Python para análisis estadístico)
TensorFlow
XGBoost

Para obtener una lista completa, consulte Built-In Model Flavors.

Sin embargo, es posible que tenga que cambiar la forma en que funciona un tipo o registrar un modelo que MLflow no admite de forma nativa. O bien, es posible que tenga que registrar un modelo que use varios elementos de varios marcos de trabajo. En estos casos, puede crear un tipo de modelo personalizado.

Para resolver el problema, MLflow ofrece el tipo PyFunc, una interfaz de modelo predeterminada para los modelos de Python. Este tipo puede registrar cualquier objeto como modelo siempre que ese objeto cumpla dos condiciones:

Implemente al menos el método predict.
El objeto de Python hereda de la mlflow.pyfunc.PythonModel clase .

Sugerencia

Los modelos serializables que implementan la API scikit-learn pueden usar el tipo scikit-learn para registrar el modelo, independientemente de si el modelo se creó con scikit-learn. Si puede conservar el modelo en formato Pickle y el objeto tiene al menos los métodos predict y predict_proba, puede usar mlflow.sklearn.log_model para registrar el modelo dentro de una ejecución de MLflow.

La manera más fácil de crear una versión para tu modelo personalizado es crear un envoltorio alrededor de tu objeto de modelo existente. MLflow serializa y empaqueta automáticamente el modelo. Los objetos de Python son serializables cuando se pueden almacenar en el sistema de archivos como un archivo, generalmente con el formato pickle. En tiempo de ejecución, el objeto se puede cargar desde ese archivo. La carga restaura todos los valores, propiedades y métodos que están disponibles cuando se guardan.

Utilice este método cuando:

Puede serializar su modelo en formato Pickle.
Desea conservar el estado del modelo justo después del entrenamiento.
Quiera personalizar cómo opera la función predict.

El código siguiente encapsula un modelo creado con XGBoost para que se comporte de forma diferente a la implementación predeterminada del tipo XGBoost. Devuelve probabilidades en lugar de clases.

from mlflow.pyfunc import PythonModel, PythonModelContext

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

    def predict(self, context: PythonModelContext, data):
        # The next line uses a prediction function. However, you could also use model.recommend(), model.forecast(), or a similar function instead.
        return self._model.predict_proba(data)

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

Use el código siguiente para registrar un modelo personalizado durante una ejecución:

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)

Sugerencia

En el código anterior, el infer_signature método usa y_probs para deducir la firma. La columna de destino contiene la clase de destino, pero el modelo devuelve dos probabilidades para cada clase.

A veces, un modelo se compone de varias partes que deben cargarse. Y, a veces, no tienes ninguna manera de serializar el modelo como un archivo Pickle. En estos casos, puede usar la PythonModel clase . Proporciona compatibilidad con una lista arbitraria de artefactos. Cada artefacto se empaqueta junto con el modelo.

Use esta técnica cuando:

No puedes serializar tu modelo en formato Pickle, a menos que tengas un mejor formato de serialización disponible.
El modelo tiene uno o varios artefactos a los que se debe hacer referencia para cargar el modelo.
Quiere conservar algunas propiedades de configuración de inferencia, por ejemplo, el número de elementos a recomendar.
Quiere personalizar la forma en que se carga el modelo y cómo funciona la predict función.

En el código siguiente se muestra cómo registrar un modelo personalizado que usa artefactos:

encoder_path = 'encoder.pkl'
joblib.dump(encoder, encoder_path)

model_path = 'xgb.model'
model.save_model(model_path)

mlflow.pyfunc.log_model("classifier", 
                        python_model=ModelWrapper(),
                        artifacts={ 
                            'encoder': encoder_path,
                            'model': model_path 
                        },
                        signature=signature)

Nota:

El modelo no se guarda en formato Pickle. En su lugar, el código guarda el modelo mediante el método save del marco que se usa.
La ModelWrapper clase se usa para encapsular el modelo, pero el modelo no se pasa como argumento al ModelWrapper constructor.
El log_model método tiene un parámetro de diccionario, artifacts. Sus claves contienen nombres de artefactos. Cada valor contiene la ruta de acceso del sistema de archivos local a un artefacto.

El contenedor de modelos correspondiente es similar al código siguiente:

from mlflow.pyfunc import PythonModel, PythonModelContext

class ModelWrapper(PythonModel):
    def load_context(self, context: PythonModelContext):
        import pickle
        from xgboost import XGBClassifier
        from sklearn.preprocessing import OrdinalEncoder
        
        self._encoder = pickle.loads(context.artifacts["encoder"])
        self._model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
        self._model.load_model(context.artifacts["model"])

    def predict(self, context: PythonModelContext, data):
        return self._model.predict_proba(data)

La rutina de entrenamiento completa es similar al código siguiente:

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

mlflow.xgboost.autolog(log_models=False)

encoder = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=np.nan)
X_train['thal'] = encoder.fit_transform(X_train['thal'].to_frame())
X_test['thal'] = encoder.transform(X_test['thal'].to_frame())

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)

encoder_path = 'encoder.pkl'
joblib.dump(encoder, encoder_path)
model_path = "xgb.model"
model.save_model(model_path)

signature = infer_signature(X, y_probs)
mlflow.pyfunc.log_model("classifier", 
                        python_model=ModelWrapper(),
                        artifacts={ 
                            'encoder': encoder_path,
                            'model': model_path 
                        },
                        signature=signature)

Un modelo a veces tiene lógica compleja o carga varios archivos de origen en el momento de la inferencia. Esta situación se produce cuando tiene una biblioteca de Python para el modelo, por ejemplo. En este escenario, debe empaquetar la biblioteca junto con el modelo para que funcionen como una sola pieza.

Use esta técnica cuando:

No puedes serializar tu modelo en formato Pickle o tienes un formato de serialización mejor disponible.
Puede almacenar los artefactos del modelo en una carpeta que almacene todos los artefactos necesarios.
El código fuente del modelo es complejo y requiere varios archivos de Python. Potencialmente, una biblioteca admite el modelo.
Quiere personalizar la forma en que se carga el modelo y cómo funciona la predict función.

MLflow admite estos tipos de modelos. Al usar MLflow, puede especificar código fuente arbitrario para empaquetar junto con el modelo, siempre y cuando el código fuente tenga un módulo de cargador. Puede especificar módulos del cargador en la llamada a log_model mediante el loader_module parámetro , que indica el espacio de nombres de Python que implementa el cargador. El code_path parámetro también es necesario. Proporciona los archivos de origen que definen el módulo del cargador. En este namespace, debe implementar una _load_pyfunc(data_path: str) función que reciba la ruta de los artefactos y devuelva un objeto que implemente al menos un predict método.

model_path = 'xgb.model'
model.save_model(model_path)

mlflow.pyfunc.log_model("classifier", 
                        data_path=model_path,
                        code_path=['src'],
                        loader_module='loader_module'
                        signature=signature)

Nota:

El modelo no se guarda en formato Pickle. En su lugar, el código guarda el modelo mediante el método save del marco que se usa.
El log_model método tiene un data_path parámetro que apunta a la carpeta que contiene los artefactos del modelo. Los artefactos pueden estar en una carpeta o un archivo. Los artefactos se empaquetan con el modelo.
El log_model método tiene un code_path parámetro que apunta a la ubicación del código fuente. El code_path valor puede ser una ruta de acceso o un único archivo. El código fuente se empaqueta con el modelo.
La _load_pyfunc función se almacena en el loader_module módulo de Python.

La carpeta src contiene el archivo loader_module.py. Ese archivo actúa como módulo del cargador y contiene las siguientes líneas:

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

    def predict(self, data):
        return self._model.predict_proba(data)

def _load_pyfunc(data_path: str):
    import os

    model = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
    model.load_model(os.path.abspath(data_path))

    return MyModel(model)

Nota:

La clase MyModel no hereda de PythonModel como se muestra anteriormente. Sin embargo, tiene una predict función .
El código fuente del modelo está en un archivo. Cualquier código fuente es adecuado. Una carpeta denominada src es ideal en esta situación.
Una _load_pyfunc función devuelve una instancia de la clase del modelo.

La rutina de entrenamiento completa es similar al código siguiente:

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)

model_path = "xgb.model"
model.save_model(model_path)

signature = infer_signature(X_test, y_probs)
mlflow.pyfunc.log_model("classifier",
                        data_path=model_path,
                        code_path=["loader_module.py"],
                        loader_module="loader_module",
                        signature=signature)

Paso siguiente

Directrices para implementar modelos de MLflow

Compartir a través de

Registro de modelos de MLflow

Prerrequisitos

¿Por qué registrar modelos en lugar de artefactos?

Uso del registro automático para registrar modelos

Modelos de registro que usan una firma, un entorno o ejemplos personalizados

Modelos de registro que usan el comportamiento de predicción modificado

Registro de modelos personalizados

Paso siguiente

Comentarios

Recursos adicionales