Customizando modelos de MLFlow com PyFuncs

Se você estuda ou trabalha com ciência de dados e aprendizado de máquina, já deve ter se deparado com alguns problemas de reproducibilidade. Tipo assim, você não consegue treinar de novo um modelo que seja tão bom quanto o que você fez 3 meses atrás. Seja porque em algum momento você usou algum hiperparâmetro diferente, ou porque você não anotou em nenhum lugar qual foi a métrica que você conseguiu com aquele treino aquele dia.

Pode ser que você também tenha tentado serializar um modelo com um framework específico que aplicava sempre o método transform para realizar predições, e aí depois você decidiu mudar de framework e o método de predição era predict e isso acabou quebrando um pouco o fluxo de servir o modelo em produção.

Se isso te parece familiar, a boa notícia é que tem uma ferramenta maravilhosa que ajuda bastante nessa gestão do seu modelo. O Mlflow se define como uma “plataforma de código aberto para o ciclo de vida de Machine Learning”.

Componentes do Mlflow

São 4 os principais componentes do Mlflow, cada um com uma responsabilidade e um benefício próprio. O legal é que você define na sua estrutura o que mais faz sentido utilizar dentre esses componentes, não sendo necessária uma adoção completa de todos.

1. Mlflow Tracking: uma ferramenta web que te auxilia a registrar os parâmetros, artefatos e métricas de cada execução de treino realizada.
2. Mlflow Projects: define como criar um ambiente virtual, com a ajuda do Anaconda e executar todos os passos (entry-points) para reproduzir a pipeline da sua execução.
3. Mlflow Models: “empacota” todos os modelos em um formato único. Também é usado para servir os modelos em produção, construindo uma imagem e subindo um servidor baseado em APIs REST.
4. Mlflow Registry: Faz a gestão das diferentes versões e estágios do seu modelo, entre Staging até Production ou Arquivado.

Mlflow Models

Existem vários artigos que te ensinam como treinar e registrar um modelo “simples” de Mlflow, usando frameworks conhecidos. Seu código fica mais ou menos assim:

import mlflow
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

with mlflow.start_run():

        mlflow.sklearn.autolog()

        data=load_iris()

        X = pd.DataFrame(data.data, columns=data.feature_names)
        y = pd.DataFrame(data.target, columns=["target"])

        clf = RandomForestClassifier(n_estimators=20)
        clf.fit(X_train, y_train)

        mlflow.end_run()

Percebe que modificamos um código “usual” de treino em 3 linhas? Quando a gente executar esse pedaço de código, o mlflow irá registrar todos os hiperparâmetros, uma série de métricas de aprendizado e também o modelo serializado como um pickle usando o formato padronizado de um modelo de MLflow. Tudo isso vai ficar armazenado no seu servidor de Mlflow Tracking e pronto para ser “empacotado” pra produção. Legal, né?

Mas essa postagem aqui não apareceu pra te ensinar a usar o Mlflow base. Parto do princípio que você pelo menos já mexeu um pouquinho e já adotou, mas precisa saber como construir um modelo customizado.

Customizar um modelo pode ser interessante se você:

• Quer que o seu método predict retorne algo além da coluna de predições
• Precisa tratar os dados com outros métodos antes de fazer as predições
• Quer servir um stack de modelos como um único modelo
• Está utilizando algum framework desconhecido pelo Mlflow

Estrutura de PyFuncs

Para construir um modelo customizado, você terá que utilizar o "sabor" pyfunc. Ele é basicamente um wrapper que pode pegar qualquer modelo genérico e transformar em um modelo de MLflow. É importante também destacar que todos os modelos de MLflow herdam a mesma classe base. Ou seja, você, ao trabalhar com PyFuncs, está estendendo as funcionalidades que já foram pré-determinadas no mlflow.sklearn, por exemplo. Mas no fim das contas, o modelo vai seguir o mesmo padrão de encapsulamento.

Isso já garante pra gente uma uniformidade em como realizar as tuas predições. Mesmo que o seu framework utilize o método transform, todos os modelos de Mlflow irão implementar um método predict por padrão.

Para construir uma PyFunc, os imports no começo do código ficam:

import mlflow
import pickle
import cloudpickle

Você vai precisar conhecer um pouco de Programação Orientada a Objetos para poder criar o seu próprio modelo, mas não é difícil. Você vai herdar a classe PythonModel e também vai precisar informar algumas coisas extras, que os frameworks padronizados pelo MLflow e o autologging abstraíam pra você antes.

Vamos começar com um exemplo de como você define um modelo:

import mlflow

class MyModelWrapper(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        with open(context.artifacts["estimator"], "rb") as f:
            self.model = pickle.load(f)

    def predict(self, context, data):
        return self.model.predict(data)

Os dois métodos acima são necessários no seu wrapper para que você consiga registrar com sucesso seu modelo. Você basicamente irá dizer para o MLflow como carregar o modelo treinado (onde ele está, qual é nome dele) e como realizar as predições.

Imagine que você precisasse carregar um preprocessador antes, chamá-lo para transformar seus dados e só depois realizar as predições. O seu código ficaria assim:

class MyModelWrapper(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        with open(context.artifacts["preprocessor"], "rb") as f:
            self.preprocessor = pickle.load(f)

        with open(context.artifacts["estimator"], "rb") as f:
            self.model = pickle.load(f)

    def predict(self, context, data):
        transformed_data = self.preprocessor.transform(data)
        return self.model.predict(transformed_data)

Percebe como é flexível? Você precisa basicamente declarar o que muda de uma estrutura padrão de modelo para a sua necessidade. Nesse caso, você pode combinar múltiplos modelos treinados em um único modelo de MLflow, ou como fiz ali em cima, um pre-processador serializado junto com o estimador. Tudo isso vai ser entendido pelo seu "pacote" de modelos de MLflow em um único predict.

Mas... pera! Tá faltando alguma coisa. Como você treina esse modelo? De onde sai esse context.artifacts? 😨

Treinando seu modelo como uma PyFunc

Bom, como você tem mais flexibilidade, você vai ter um trabalinho maior também pra definir seu treino. Apesar de não ser estritamente necessário, recomendo fortemente que você inclua na sua classe MyModelWrapper o médoto fit. Aí você vai conseguir empacotar tudo em um fluxo de trabalho só, e vamos mostrar um exemplo de como fazer. Perceba que no nosso construtor estamos incluindo o preprocessador e o estimador:

class MyModelWrapper(mlflow.pyfunc.PythonModel):

    def __init__(self):

        self.preprocessor = MinMaxScaler()
        self.model = SVC()

    def load_context(self, context):
        with open(context.artifacts["preprocessor"], "rb") as f:
            self.processor = pickle.load(f)
        with open(context.artifacts["estimator"], "rb") as f:
            self.estimator = pickle.load(f)

    def fit(self, data:pd.DataFrame) -> None:
        # Making the splits
        X = data.drop('target', axis=1)
        y = data['target']

        # Training the transformer and scaling the data
        X_scaled = self.preprocessor.fit_transform(data)

        X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.25)

        # Training your model Pipeline
        self.model.fit(X_train, y_train)

        # Evaluating and logging metric with MLFlow
        y_pred = self.model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric("accuracy", accuracy)

        # Dumping the fitted objects
        with open("fitted_processor.pkl", "wb") as f:
            cloudpickle.dump(self.preprocessor, f)

        with open("fitted_model.pkl", "wb") as f:
            cloudpickle.dump(self.model, f)

    def predict(self, context, data):
        transformed_data = self.preprocessor.transform(data)
        return self.model.predict(transformed_data)

Bom, primeira coisa que eu gostaria de pontuar é que estou utilizando o cloudpickle para persistir os modelos treinados. O MLflow se dá muito bem com esse serializador, e depois você pode usar o pickle para desserializar numa boa. Mas, caso necessário, você pode utilizar o serializador de sua escolha, desde que você ajuste depois o método load_context de acordo. Além disso, mesmo antes de termos definido a nossa execução, já podemos dizer qual métrica queremos acompanhar com o MLFlow.

Pra treinar o algoritmo, você vai precisar definir:

• O ambiente de Anaconda para replicar seu experimento
• Definir a classe do seu modelo
• Definir um dicionário de artefatos
• Definir o nome do diretório que contém os artefatos

Vamos ver o código então:

with mlflow.start_run():

    model = MyModelWrapper()

    data = read_data()

    model.fit(data)

    conda_env = "conda.yaml"

    artifacts = {
        "preprocessor" : "fitted_processor.pkl",
        "estimator" : "fitted_model.pkl"
    }

    mlflow.pyfunc.log_model(
        artifact_path='model',
        artifacts=artifacts,
        python_model = model,
        conda_env = conda_env
    )

    mlflow.end_run()

Esse conda.yaml é um arquivo que precisa estar na sua estrutura de execução, na mesma organização de pastas em que você executará o seu módulo de treino.

Além disso, o método fit faz o dump dos objetos serializados no seu diretório local. Então você abre um dicionário chamado artifacts e dá o nome das chaves de cada artefato que você está logando. Perceba que no load_context pegamos justamente os valores das chaves preprocessor e estimator do dicionário artifact. Você também pode adicionar a esse dicionário um plot de gráfico em JPEG ou até mesmo um arquivo CSV com as perdas em cada época da sua rede neural, por exemplo.

Enfim, artefatos servem para que você carregue informação junto com a sua execução de treino lá pro servidor do Mlflow Tracking. Ou mesmo modelos treinados, que precisarão aparecer no seu load_context depois 😄

Fazendo predições

Como comentei mais cedo, todos os modelos que forem registrados como modelos de MLflow terão a mesma estrutura. Seja ele um modelo “padrão”, utilizando frameworks conhecidos, ou mesmo o customizado, como acabamos de ver.

Pra fazer as predições, pra qualquer modelo de MLflow, a gente vai precisar fazer o seguinte:

import mlflow
import pandas as pd

logged_model = 'runs:/<run_id>/<artifact_directory>'
loaded_model = mlflow.pyfunc.load_model(logged_model)

loaded_model.predict(pd.DataFrame(data))

Então, só trocando a run_id você já pode usar o mesmo método de predição. Percebe o valor disso? O que a gente fez foi basicamente editar e estender tanto o que carregar quando eu chamar mlflow.pyfunc.load_model, como também como exatamente realizar as predições, chamando o método padronizado predict.

Parabéns, você agora já sabe como montar um modelo customizado com uma ferramenta mega poderosa 🥳

O que mais?

Se você se interessou por esse tópico e quer explorar ainda mais o MLflow Models, não deixe de conferir a documentação das PyFuncs, pois isso poderá te ajudar bastante no seu desenvolvimento.

De verdade, espero que esse tutorial seja útil pra você.

Tem dúvidas? Deixe seu comentário 😃