Uso de modelos particionados

Muitos conjuntos de dados podem ser divididos em vários subconjuntos independentes. Por exemplo, um conjunto de dados com dados de vendas de uma rede de lojas pode ser particionado por número de loja. Um modelo separado pode então ser treinado para cada partição. As operações de treinamento e inferência nas partições podem ser paralelizadas, reduzindo o tempo em hora local nessas operações. Além disso, como as lojas individuais provavelmente diferem significativamente na forma como seus recursos afetam suas vendas, essa abordagem pode levar a uma inferência mais precisa no nível da loja.

O Snowflake Model Registry oferece suporte ao processamento distribuído de treinamento e inferência de dados particionados quando:

  • O conjunto de dados contém uma coluna que identifica de forma confiável as partições nos dados.

  • Os dados em cada partição individual não estão correlacionados com os dados nas outras partições e contêm linhas suficientes para treinar o modelo.

Os modelos podem ser sem estado (o treinamento é realizado toda vez que a inferência é chamada) ou com estado (o treinamento é realizado uma vez antes da inferência e retido para uso em várias operações de inferência).

Com o Snowflake Model Registry, implemente treinamento e inferência particionados usando modelos personalizados. Durante a inferência, o método de inferência de modelo particiona o conjunto de dados, gera previsões para cada partição em paralelo usando todos os nós e núcleos em seu warehouse e combina os resultados em um único conjunto de dados posteriormente.

Nota

Para modelos particionados, é importante distinguir o modelo registrado dos modelos individuais que são criados por ou compõem o modelo registrado. Sempre que possível, vamos nos referir aos modelos subjacentes individuais como submodelos.

Nota

O treinamento e a inferência particionados exigem o Snowpark ML (pacote snowflake-ml-python) versão 1.5.0 ou posterior.

Definição e registro do modelo

A classe do modelo particionado herda de snowflake.ml.model.custom_model.CustomModel, e os métodos de inferência são declarados com o decorador @custom_model.partitioned_inference_api (Snowpark ML versão 1.5.4 ou posterior) ou @custom_model.inference_api (Snowpark ML versão 1.5.0 a 1.5.3). Consulte Traga seus próprios tipos de modelos por meio de arquivos serializados para obter informações sobre a definição de modelos personalizados padrão.

import pandas as pd

from snowflake.ml.model import custom_model

class ExamplePartitionedModel(custom_model.CustomModel):

  @custom_model.partitioned_inference_api
  def predict(self, input: pd.DataFrame) -> pd.DataFrame:
      # All data in the partition will be loaded in the input dataframe.
      #… implement model logic here …
      return output_df

my_model = ExamplePartitionedModel()
Copy

Ao registrar o modelo, forneça um function_type de TABLE_FUNCTION no dicionário options junto com quaisquer outras opções que seu modelo exija.

from snowflake.ml.registry import Registry

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
model_version = reg.log_model(my_model,
  model_name="my_model",
  version_name="v1",
  options={"function_type": "TABLE_FUNCTION"},    ###
  conda_dependencies=["scikit-learn"],
  sample_input_data=train_features
)
Copy

Se o seu modelo particionado também tiver funções regulares (não de tabela) como métodos, você poderá usar o dicionário method_options para especificar o tipo de cada método.

model_version = reg.log_model(my_model,
    model_name="my_model",
    version_name="v1",
    options={
      "method_options": {                                 ###
        "METHOD1": {"function_type": "TABLE_FUNCTION"},   ###
        "METHOD2": {"function_type": "FUNCTION"}          ###
      }
    }
    conda_dependencies=["scikit-learn"],
    sample_input_data=train_features
)
Copy

Inferência de modelo particionado

Use o método run de um objeto Python ModelVersion para invocar os métodos da função de tabela de forma particionada, passando partition_column para especificar o nome da coluna com um valor numérico ou de cadeia de caracteres que identifica a partição de cada registro. Como de costume, você pode passar um DataFrame Snowpark ou pandas (este último é útil para testes locais). Você receberá o mesmo tipo de DataFrame como resultado. Nesses exemplos, a inferência é particionada em um número de loja.

model_version.run(
  input_df,
  function_name="PREDICT",
  partition_column="STORE_NUMBER"
)
Copy

Você também pode chamar as funções da tabela de modelos diretamente com SQL, como mostrado aqui.

SELECT output1, output2, partition_column
  FROM input_table,
      TABLE(
          my_model!predict(input_table.input1, input_table.input2)
          OVER (PARTITION BY input_table.store_number)
      )
  ORDER BY input_table.store_number;
Copy

Os dados de entrada são divididos automaticamente entre os nós e núcleos no warehouse, e as partições são processadas em paralelo.

Para obter mais informações sobre a sintaxe da função de tabela, consulte Como chamar uma UDF com o SQL.

Modelos particionados sem estado

Na aplicação mais simples de modelos particionados, o treinamento e a inferência são feitos quando o predict é chamado. O modelo é ajustado, a inferência é executada e o modelo ajustado é descartado imediatamente depois. Esse tipo de modelo é chamado de «sem estado» porque nenhum estado de ajuste é armazenado. Aqui está um exemplo em que cada partição treina um modelo XGBoost:

class ExampleStatelessPartitionedModel(custom_model.CustomModel):

  @custom_model.partitioned_inference_api
  def predict(self, input_df: pd.DataFrame) -> pd.DataFrame:
      import xgboost
      # All data in the partition will be loaded in the input dataframe.
      # Construct training data by transforming input_df.
      training_data = # ...

      # Train the model.
      my_model = xgboost.XGBRegressor()
      my_model.fit(training_data)

      # Generate predictions.
      output_df = my_model.predict(...)

      return output_df

my_model = ExampleStatelessPartitionedModel()
Copy

Consulte o Guia de início rápido do modelo personalizado particionado para obter um exemplo de um modelo particionado sem estado, incluindo dados de amostra.

Modelos particionados com estado

Também é possível implementar modelos particionados com estado que carregam o estado de ajuste do submodelo armazenado. Você faz isso fornecendo modelos na memória por meio do snowflake.ml.model.custom_model.ModelContext ou fornecendo caminhos de arquivos que apontam para artefatos de modelos ajustados e carregando-os durante a inferência.

O exemplo a seguir mostra como fornecer modelos na memória para o contexto do modelo.

from snowflake.ml.model import custom_model

# `models` is a dict with model ids as keys, and fitted xgboost models as values.
models = {
  "model1": models[0],
  "model2": models[1],
  ...
}

model_context = custom_model.ModelContext(
  models=models
)
my_stateful_model = MyStatefulCustomModel(model_context=model_context)
Copy

Ao registrar my_stateful_model, os submodelos fornecidos no contexto são armazenados junto com todos os arquivos de modelo. Eles podem então ser acessados na lógica do método de inferência, recuperando-os do contexto, conforme mostrado abaixo:

class ExampleStatefulModel(custom_model.CustomModel):

  @custom_model.partitioned_inference_api
  def predict(self, input: pd.DataFrame) -> pd.DataFrame:
    model1 = self.context.model_ref("model1")
    # ... use model1 for inference
Copy

Também é possível acessar os modelos de forma programática por meio do ID da partição no método predict. Se uma coluna de partição for fornecida como um recurso de entrada, ela poderá ser usada para acessar um modelo ajustado para a partição. Por exemplo, se a coluna de partição for MY_PARTITION_COLUMN, a seguinte classe de modelo poderá ser definida:

class ExampleStatefulModel(custom_model.CustomModel):

  @custom_model.partitioned_inference_api
  def predict(self, input: pd.DataFrame) -> pd.DataFrame:
    model_id = input["MY_PARTITION_COLUMN"][0]
    model = self.context.model_ref(model_id)
    # ... use model for inference
Copy

Da mesma forma, os submodelos podem ser armazenados como artefatos e carregados em tempo de execução. Essa abordagem é útil quando os modelos são muito grandes para caber na memória. Fornece caminhos de arquivo de cadeia de caracteres para o contexto do modelo. Os caminhos dos arquivos podem ser acessados durante a inferência com self.context.path(artifact_id). Para obter mais informações, consulte Definição do contexto do modelo por argumentos de palavras-chave.

Exemplo

Consulte o Guia de início rápido do modelo particionado para ver um exemplo, incluindo dados de amostra.

Consulte o Guia de início rápido de Inferência de vários modelos no Snowflake para ver um exemplo de modelo personalizado particionado com estado.