分割モデルの使用

多くのデータセットは、複数の独立したサブセットに分割することができます。例えば、チェーン店の売上データを含むデータセットは、店舗番号でパーティショニングできます。各パーティションごとに別々のモデルをトレーニングすることができます。パーティションに対するトレーニングと推論操作は並列化でき、これらの操作にかかるウォールクロック時間を短縮できます。さらに、個々のストアの特徴が売上に与える影響は大きく異なる可能性が高いため、このアプローチはストアレベルでより正確な推論につながります。

次の場合、Snowflake Model Registryは、パーティショニングされたデータのトレーニングと推論の分散処理に対応します。

  • データセットには、データのパーティションを確実に特定する列が含まれています。

  • 各パーティションのデータは、他のパーティションのデータと相関がなく、モデルを訓練するのに十分な行数を含んでいます。

モデルはステートレス（推論が呼び出されるたびに学習が実行される）である場合もあれば、ステートフル（推論の前に一度学習が実行され、複数の推論操作で使用するために保持される）である場合もあります。

Snowflake Model Registry を使用して、 カスタムモデル を使用して、パーティショニングされた学習と推論を実装します。推論中、モデル推論メソッドはデータセットを分割し、ウェアハウス内のすべてのノードとコアを使用して各パーティションの予測を並列に生成し、その結果を後で単一のデータセットに結合します。

注釈

パーティショニングされたモデルの場合、登録されたモデルを、登録されたモデルによって作成された、あるいは登録されたモデルを構成する個々のモデルから区別することが重要です。可能な限り、個々の基になるモデルをサブモデルとしてリファレンスします。

注釈

パーティショニングされたトレーニングと推論にはSnowpark ML （snowflake-ml-python package）version 1.5.0以降が必要です。

モデルの定義とログ記録

パーティションモデルクラスは snowflake.ml.model.custom_model を継承しています。CustomModel、推論メソッドは @custom_model.partitioned_inference_api デコレータ（Snowpark ML バージョン1.5.4以降）または @custom_model.inference_api デコレータ（Snowpark ML バージョン1.5.0〜1.5.3）で宣言されます。標準カスタムモデルの定義に関する情報については、 シリアライズされたファイルを介して独自のモデルタイプを持ち込むことができます。 を参照してください。

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

モデルをログに記録する際には、 options 辞書内の TABLE_FUNCTION の function_type を、モデルが必要とするその他の オプション とともに提供します。

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

パーティショニングされたモデルがメソッドとして通常の(テーブルではない)関数も持つ場合、代わりに method_options ディクショナリを使用して各メソッドのタイプを指定することができます。

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

分割モデル推論

Pythonの ModelVersion オブジェクトの run メソッドを使用して、パーティショニングされた方法でテーブル関数のメソッドを呼び出します。 partition_column を渡して、各記録のパーティションを識別する数値または文字列値を含む列の名前を指定します。いつものように、Snowparkまたはpandas DataFrame を渡すことができます（後者はローカルテストに便利）。その結果、同じタイプの DataFrame を受けとります。これらの例では、推論はストア番号で分割されています。

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

ここに示すように、 SQL を使ってモデルテーブル関数を直接呼び出すこともできます。

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

入力データは自動的にウェアハウス内のノードとコアに分割され、パーティションは並列処理されます。

テーブル関数の構文についての情報は、 SQL で UDF を呼び出します。 をご参照ください。

ステートレス分割モデル

パーティショニング・モデルの最も単純なアプリケーションでは、 predict が呼び出されると、学習と推論の両方が行われます。モデルがフィッティングされ、推論が実行され、フィッティングされたモデルは直後に破棄されます。このタイプのモデルは、フィット状態が保存されないため、「ステートレス」と呼ばれます。各パーティションが 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

サンプルデータを含むステートレス・パーティションド・モデルの例については、 パーティションド・モデル・クイックスタート・ガイド を参照してください。

ステートフル・パーティションド・モデル

また、保存されたサブモデルのフィット状態をロードするステートフルなパーティショニング・モデルを実装することも可能です。これは、 snowflake.ml.model.custom_model.ModelContext を介してメモリ上にモデルをプロバイダーするか、フィッティングされたモデルアーティファクトを指すファイルパスを提供し、推論中にそれらをロードすることによって行います。

次の例は、メモリ上のモデルをモデルコンテキストにプロバイダーする方法を示しています。

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

my_stateful_model をログに記録する場合、コンテキストでプロバイダーされたサブモデルは、すべてのモデルファイルとともに保存されます。そして、推論メソッドのロジックでは、以下のようにコンテキストから取得することでアクセスできます。

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

また、 predict メソッド内のパーティション ID によって、プログラム的にモデルにアクセスすることも可能です。パーティション列が入力特徴としてプロバイダーされた場合、そのパーティションに適合したモデルにアクセスするために使用できます。たとえば、パーティション列が MY_PARTITION_COLUMN の場合、次のようなモデルクラスが定義できます。

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

同様に、サブモデルを成果物として保存し、ランタイムにロードすることもできます。この方法は、モデルが大容量でメモリに収まらない場合に有効です。モデルコンテキストの文字列ファイルパスを指定します。ファイルパスは推論中に self.context.path(artifact_id) でアクセス可能。詳細については、 キーワード引数によるモデルコンテキストの定義 をご参照ください。

例

サンプルデータを含む例については、 Partitioned Model Quickstart Guide を参照してください。

ステートフルパーティショニングカスタムモデルの例については、 Many Model Inference in Snowflake Quickstart Guide を参照してください。

言語: 日本語