Snowflake Model Registry: modelos personalizados particionados

Muitos conjuntos de dados podem ser facilmente particionados 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 cada loja provavelmente difere um pouco na forma como seus recursos afetam suas vendas, essa abordagem pode realmente levar a inferências mais precisas 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.

  • O modelo não tem estado: ele executa tanto o ajuste (treinamento) quanto a inferência (previsão) sempre que é chamado e não persiste pesos ou outros estados do modelo entre as chamadas.

Com o Snowflake Model Registry, você implementa treinamento e inferência particionados usando modelos personalizados. Ao usar o modelo, o registro particiona o conjunto de dados, ajusta e prevê as partições em paralelo usando todos os nós e núcleos no warehouse e, posteriormente, combina os resultados em um único conjunto de dados.

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 personalizado

Conforme explicado em Como escrever a classe de modelo personalizado, você declara métodos de inferência de modelo personalizados com o decorador @custom_model.partitioned_inference_api (Snowpark ML versão 1.5.4 ou posterior) ou decorador @custom_model.inference_api (Snowpark ML versão 1.5.0 a 1.5.3).

class ExampleForecastingModel(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 = ExampleForecastingModel()
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.

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = 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 modelo personalizado 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

Realização de treinamento e inferência

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. Nestes exemplos, a partição é feita no número de uma loja.

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

Também é possível chamar esses métodos usando dados particionados de SQL, conforme 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.

Dica

Muitos conjuntos de dados podem ser particionados de mais de uma maneira. Como a coluna de partição é especificada quando você chama o modelo, não quando você o registra, você pode facilmente experimentar diferentes esquemas de particionamento sem alterar o modelo. Por exemplo, no conjunto de dados hipotético de vendas em lojas, é possível particionar por número de loja ou por estado ou província para ver qual prevê com mais eficácia.

Isso também significa que você não precisa de um modelo separado para processamento não particionado. Se você não especificar uma coluna de partição, nenhum particionamento será feito e todos os dados serão processados juntos, como de costume.

Exemplo

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