Treinamento de modelos em partições de dados

Uso de treinamento de vários modelos (MMT) para treinar vários modelos de aprendizado de máquina de forma eficiente em partições de dados. Ele lida automaticamente com orquestração distribuída, armazenamento de modelo e persistência de artefatos.

O MMT particiona seu Snowpark DataFrame por uma coluna especificada e treina modelos separados em cada partição em paralelo. Concentre-se na lógica de treinamento do seu modelo enquanto o MMT lida com a complexidade da infraestrutura e é dimensionado automaticamente.

Você pode usar o MMT para treinar vários modelos de forma eficiente em diferentes segmentos de dados. Essa ferramenta é ideal para cenários como o treinamento de modelos de previsão de vendas específicos da região, a criação de sistemas de recomendação personalizada em que cada grupo de clientes requer seu próprio modelo ou a criação de modelos preditivos específicos do segmento. O MMT lida automaticamente com o treinamento do modelo distribuído, eliminando a complexidade do gerenciamento da infraestrutura de computação distribuída.

Você pode usar o MMT para treinar modelos usando modelos e estruturas de aprendizado de máquina de código aberto, como XGBoost, scikit-learn, PyTorch e TensorFlow. O MMT serializa automaticamente artefatos do modelo, para que você possa acessá-los no momento da inferência.

Você também pode implementar a interface ModelSerde para treinar modelos personalizados ou usar estruturas ML sem suporte. Isso permite integrar o MMT com qualquer estrutura de aprendizado de máquina ou arquitetura de modelo personalizada que você use.

Importante

Antes de usar o MMT, certifique-se de ter o seguinte:

  • Ambiente de tempo de execução do contêiner: o MMT requer um ambiente de tempo de execução do contêiner do Snowflake ML.

  • Permissões de acesso a estágio: o MMT armazena automaticamente artefatos de modelo em estágios do Snowflake. Certifique-se de ter as permissões apropriadas para acessar o estágio nomeado especificado.

  • Suporte a estruturas de ML: as integrações internas estão disponíveis para XGBoost, scikit-learn, PyTorch e TensorFlow. Para modelos personalizados, implemente a interface ModelSerde.

A seção a seguir mostra como usar o MMT em um fluxo de trabalho de exemplo.

Treinamento de um modelo com o MMT

Esta seção demonstra o fluxo de trabalho completo do MMT em cinco etapas principais:

  1. Importar seus dados - carregar dados de treinamento usando o Snowpark

  2. Definir a função de treinamento - definir a função de treinamento

  3. Treinar modelos em partições - usar o MMT para treinar modelos em cada partição em paralelo

  4. Acessar modelos treinados - recuperar e usar os modelos treinados para cada partição

  5. Persistência e recuperação de modelo - salvar modelos em estágios e restaurá-los posteriormente

O fluxo de trabalho lida automaticamente com treinamento distribuído, serialização de modelo e armazenamento de artefatos em suas partições de dados.

Importação de dados

Use uma sessão de Snowpark para começar a importar seus dados. A função de treinamento de vários modelos divide os dados importados em diferentes partições usando a coluna especificada.

Antes de usar o MMT, crie uma sessão do Snowpark. Para obter mais informações, Como criar uma sessão para o Snowpark Python.

O código a seguir usa uma sessão do Snowpark para importar seus dados de treinamento.

# Example: sales_data with columns: region, feature1, feature2, feature3, target
sales_data = session.table("SALES_TRAINING_DATA")
Copy

Definição da função de treinamento

Depois de obter seus dados, você define a função de treinamento que o MMT usa para treinar modelos entre partições. A função de treinamento recebe um conector de dados e um objeto de contexto que o aponta para a partição de dados na qual é treinado. Esta seção tem exemplos que definem uma função de treinamento para treinar um modelo XGBoost além de exemplos que aproveitam TensorFlow e PyTorch.

Sua função de treinamento deve ter exatamente esta assinatura: (data_connector, context). Para cada partição de dados, chamadas train_xgboost_model do MMT com os seguintes argumentos:

  • data_connector: um conector de dados que fornece acesso aos dados que o MMT particiona. train_xgboost_model converte esse dataframe em pandas.

  • context: um objeto que fornece partition_id à função train_xgboost_model. Este ID é o nome da coluna em que você está particionando.

Você não chama esta função sozinho. O MMT trata da execução em todas as partições.

Use o código a seguir para definir sua função de treinamento. Depois de alterar o código para refletir os recursos em seus dados, você poderá passá-lo para a função MMT.

Use XGBoost para treinar modelos em partições de dados. XGBoost oferece excelente desempenho para dados estruturados e trata valores ausentes automaticamente.

def train_xgboost_model(data_connector, context):
    df = data_connector.to_pandas()
    print(f"Training model for partition: {context.partition_id}")

    # Prepare features and target
    X = df[['feature1', 'feature2', 'feature3']]
    y = df['target']

    # Train the model
    from xgboost import XGBRegressor
    model = XGBRegressor(
        n_estimators=100,
        max_depth=6,
        learning_rate=0.1,
        random_state=42
    )
    model.fit(X, y)
    return model

trainer = ManyModelTraining(train_xgboost_model, "model_stage")
Copy

Treinamento de modelos em partições

Depois de definir sua função de treinamento, você pode usar o MMT para treinar modelos em partições. Especifique a coluna pela qual particionar e o estágio onde os modelos são salvos.

O código a seguir particiona os dados pela coluna region e usa a função train_xgboost_model para treinar modelos separados para cada região em paralelo.

Por exemplo, se os seguintes fossem os valores possíveis para a coluna region:

  • Norte

  • Sul

  • Leste

  • Oeste

  • Centro

A função ManyModelTraining criaria uma partição de dados separada para cada uma das regiões anteriores e treinaria um modelo em cada partição.

from snowflake.ml.modeling.distributors.many_model import ManyModelTraining

trainer = ManyModelTraining(train_xgboost_model, "model_stage") # Specify the stage to store the models
training_run = trainer.run(
    partition_by="region",  # Train separate models for each region
    snowpark_dataframe=sales_data,
    run_id="regional_models_v1" # Specify a unique ID for the training run
)

# Monitor training progress
final_status = training_run.wait()
print(f"Training completed with status: {final_status}")
Copy

Os modelos são armazenados no estágio em run_id/{partition_id}, onde partition_id é o valor da coluna de partição.

Acesso a modelos treinados

Depois que o MMT termina, você tem modelos treinados para cada partição de dados armazenada no estágio especificado. Cada modelo é treinado em dados específicos de sua partição. Por exemplo, um modelo “Norte” é treinado apenas em dados da região Norte.

O objeto de execução de treinamento fornece métodos para acessar esses modelos e verificar o status de treinamento de cada partição.

O código a seguir recupera o verifica o status da execução de treinamento e recupera os modelos treinados para cada partição:

if final_status == RunStatus.SUCCESS:
    # Access models for each partition
    for partition_id in training_run.partition_details:
        trained_model = training_run.get_model(partition_id)
        print(f"Model for {partition_id}: {trained_model}")

        # You can now use the model for predictions or further analysis
        # Example: model.predict(new_data)
else:
    # Handle training failures
    for partition_id, details in training_run.partition_details.items():
        if details.status != "DONE":
            print(f"Training failed for {partition_id}")
            error_logs = details.logs
Copy

Persistência e recuperação do modelo

O MMT persiste automaticamente modelos treinados para o estágio especificado do Snowflake durante o processo de treinamento. Cada modelo é armazenado com um caminho estruturado que inclui o ID da execução e o identificador da partição, facilitando a organização e a recuperação de modelos posteriormente.

A persistência automática significa que você não precisa salvar os modelos manualmente. O MMT cuida da serialização e do armazenamento para você, eliminando o risco de perda de modelos treinados devido a tempos limite de sessão ou problemas de conexão.

Você pode restaurar execuções de treinamento anteriores e acessar seus modelos, mesmo após o término da sessão original. Esse mecanismo de persistência permite:

  • Retomar o trabalho em diferentes sessões

  • Compartilhar modelos treinados com os membros da equipe

  • Criar fluxos de trabalho de controle de versão de modelo

  • Integrar com pipelines de inferência downstream

Os modelos são salvos automaticamente no estágio especificado e podem ser recuperados posteriormente:

# Restore training run from stage
restored_run = ManyModelTraining.restore_from("regional_models_v1", "model_stage")

# Access models from restored run
north_model = restored_run.get_model("North")
south_model = restored_run.get_model("South")
Copy

Treinamento de modelos personalizados

Para modelos personalizados ou estruturas de ML sem suporte, implemente a interface ModelSerde. Você pode definir sua própria lógica de serialização e desserialização para modelos personalizados. Isso permite integrar o MMT com qualquer estrutura de aprendizado de máquina ou arquitetura de modelo personalizada que você use.

from snowflake.ml.modeling.distributors.many_model import ModelSerde

class CustomModelSerde(ModelSerde):
    def serialize(self, model, path):
        # Custom serialization logic
        pass

    def deserialize(self, path):
        # Custom deserialization logic
        pass

def train_custom_model(data_connector, context):
    # Your custom training logic
    model = your_custom_model_training(data_connector.to_pandas())
    return model

trainer = ManyModelTraining(
    train_custom_model,
    "custom_model_stage",
    model_serde=CustomModelSerde()
)
Copy

Integração com o Registro de modelos

O MMT pode ser integrado ao Registro de modelos do Snowflake para aprimoramento de gerenciamento de modelos. O Registro de modelos fornece controle de versão de modelo centralizado, rastreamento de metadados e gerenciamento de implantação em toda a sua organização. Essa integração é particularmente valiosa ao treinar vários modelos com o MMT, pois ajuda você a organizar, rastrear e controlar todos os modelos específicos de partição a partir de um único local.

Usar o Registro de modelos com o MMT permite que você faça o seguinte:

  • Rastrear diferentes iterações de seus modelos específicos de partição

  • Armazenar métricas de desempenho de modelos, parâmetros de treinamento e informações de linhagem

  • Gerenciar quais versões do modelo são implantadas em produção para cada partição

  • Compartilhar modelos entre equipes com controles de acesso e documentação adequados

  • Implementar fluxos de trabalho de aprovação e rastreamento de conformidade para implantações de modelo

# Register trained models to Model Registry
for partition_id in training_run.partition_details:
    model = training_run.get_model(partition_id)

    # Register to Model Registry
    model_ref = registry.log_model(
        model,
        model_name=f"sales_model_{partition_id.lower()}",
        version_name="v1"
    )
Copy