Snowflake Model Registry: Partitioned Custom Models

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

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

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

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

  • このモデルはステートレスであり、呼び出されるたびにフィッティング(トレーニング)と推論(予測)の両方を実行し、呼び出しの間に重みなどのモデルの状態を保持することはありません。

Snowflake Model Registry を使用すると、 カスタムモデル を使用して、パーティショニングされたトレーニングと推論を実装できます。このモデルを使用する場合、レジストリはデータセットをパーティショニングし、ウェアハウス内のすべてのノードとコアを使用して並列にパーティショニングをフィッティングして予測し、その後に結果を1つのデータセットにまとめます。

注釈

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

カスタムモデルの定義とログ

カスタムモデルクラスの書き込み で説明したように、 @custom_model.partitioned_inference_api デコレーター(Snowpark ML バージョン 1.5.4 以降)または @custom_model.inference_api デコレーター(Snowpark ML バージョン 1.5.0 から 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

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

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

カスタムモデルがメソッドとして通常の(テーブルではない)関数も持っている場合、代わりに 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 を受けとります。これらの例では、店舗番号でパーティションしています。

mv.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

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

Tip

多くのデータセットは複数の方法でパーティショニングできます。パーティション列はモデルをログに記録するときではなく、モデルを呼び出すときに指定されるので、モデルを変更することなく、異なるパーティショニングスキームを簡単に試すことができます。例えば、仮想の店舗売上データセットでは、店舗番号でパーティショニングするか、州や県でパーティショニングして、どちらがより効果的に予測できるかを調べることができます。

これは、パーティショニングされない処理用に別のモデルを用意する必要がないことも意味します。パーティション列を指定しない場合、パーティショニングは行われず、通常通りすべてのデータがまとめて処理されます。

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