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.

Definição e registro do modelo

The partitioned model class inherits from snowflake.ml.model.custom_model.CustomModel, and inference methods are declared with the @custom_model.partitioned_api decorator. See Traga seus próprios tipos de modelos por meio de arquivos serializados for information on defining standard custom models.

import pandas as pd

from snowflake.ml.model import custom_model

class ExamplePartitionedModel(custom_model.CustomModel):

  @custom_model.partitioned_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 the run method of a Python ModelVersion object to invoke the table function methods in a partitioned fashion, passing partition_column to specify the name of the column that contains a numeric or string value that identifies the partition of each record. As usual, you may pass a Snowpark or pandas DataFrame (the latter is useful for local testing). You will receive the same type of DataFrame as the result. In these examples, inference is partitioned on a store number.

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 predict:

class ExampleStatelessPartitionedModel(custom_model.CustomModel):

  @custom_model.partitioned_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.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.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.