Snowflakeモデルレジストリへのカスタムモデルの保存

Snowflakeモデルレジストリ により、(ログ)モデルを登録し、Snowflake内で推論に使用することができます。レジストリはいくつかのタイプのモデルをサポートしています。

  • Snowpark ML Modeling

  • scikit-learn

  • XGBoost

  • LightGBM

  • CatBoost

  • PyTorch

  • TensorFlow

  • MLFlow PyFunc

  • センテンストランスフォーマー

  • Hugging Faceパイプライン

モデルレジストリ API では、 snowflake.ml.model.custom_model.CustomModel クラスから派生したシリアル化可能なモデルであれば、外部ツールを使って学習したモデルやオープンソースリポジトリから取得したモデルなど、他のタイプのモデルもログに記録することができます。

注釈

このガイド は、カスタムモデルのログの例を提供します。

AWS SageMaker またはAzure ML からのモデルのインポートについては、 このブログ投稿 および付属のクイックスタートもご参照ください。

このトピックでは、モデルを作成し、Snowflakeモデルレジストリにログを記録し、展開する方法を説明します。この機能の一般的なユースケースは、予測器や分類器に続くいくつかの変換器や推測器のような、複数のクラスのパイプラインを定義することです。このような場合、カスタムモデルクラス自体は比較的単純であり、これらのクラスを順番に呼び出し、その結果を次のクラスの入力として渡します。

モデルコンテキストの定義

モデルは多くの場合、コードに加えて、設定ファイルやモデルの重みのような1つ以上の静的ファイルを必要とします。カスタムモデルは、レジストリがモデルと一緒にそれらをパックすることがわかるように、そのようなアーティファクトの成果物に関する情報を提供する必要があります。アーティファクトは、 ModelContext クラスを使って以下のように宣言できます。

from snowflake.ml.model import custom_model

mc = custom_model.ModelContext(
    artifacts={
        'config': 'local_model_dir/config.json'
    },
)
Copy

モデルコンテキストのアーティファクトファイルへのパスは、モデルがログに記録される環境の現行ディレクトリからの相対パスです。Snowflake Model Registryはこの情報を使用して、必要なすべてのコードとデータがモデルを実行するウェアハウスに展開されるようにします。ランタイム時に、モデルは self.context.path('config') を呼び出すことで、それらのアーティファクトを見つけることができます(値 'config' は、 ModelContext に渡されたディクショナリのキーと同じです)。

静的ファイルの他に、モデルはサポートされているタイプの他のモデルやパイプライン(例: snowflake.ml.modeling.pipeline やscikit-learnモデル)を構成することができます。レジストリはすでにこれらのタイプのオブジェクトをログに記録する方法を知っているため、 model_refs を使ってモデルコンテキストに直接Pythonオブジェクトを渡すことができます。これらのオブジェクトを自分でパッケージする必要はありません。これは、ブートストラップ集計(バギング)や前処理や後処理に役立ちます。

注釈

model1 および model2 は、レジストリによってネイティブにサポートされた任意のタイプのモデルのオブジェクトです。 feature_preprocsnowflake.ml.modeling.pipeline オブジェクトです。

mc = custom_model.ModelContext(
    artifacts={
        'config': 'local_model_dir/config.json'
    },
    models={
        'm1': model1,
        'm2': model2,
        'feature_preproc': preproc
    }
)
Copy

モデルレジストリは、モデルをログに記録する際にこれらのモデル参照をシリアル化し、実行時にそれらを再水和します。モデルは、例えば self.context.model_ref('m1') を使って下位モデルへの参照を取得することができます。 self.context.model_ref('m1').predict() モデルが predict メソッドを公開している場合、コードは取得したモデル参照からそのメソッドを直接呼び出すことができます。

要約すると、カスタムモデルのコンテキストは、モデルのアーティファクト(モデルによって使用され、コードと一緒にSnowflakeに保存されなければならないローカルファイル)とともにシリアル化されるPythonオブジェクトを宣言します。モデルクラスは、コードとアーティファクトを見つけるためにコンテキストを使用します。これは、モデルが実行されているのがローカルでもSnowflakeでも機能します。

カスタムモデルクラスの書き込み

カスタムモデルのログと展開方法をモデルレジストリに伝えるには、 snowflake.ml.model.custom_model.CustomModel を継承します。

モデルは複数の推論メソッドを公開することができます(たとえば、scikit-learnモデルは predictpredict_proba メソッドを公開します)。カスタムモデルで推論関数を宣言するには、それらをサブクラスのパブリックメソッドとして定義し、 @custom_model.inference_api で装飾します。このデコレーターは、メソッドがモデルのpublic API の一部であることを示し、Pythonや SQL からモデルレジストリ経由で呼び出せるようにします。 inference_api で装飾されたメソッドは、pandas DataFrame を受け取り、返す必要があります。メソッドはいくつでも inference_api で装飾できます。

注釈

パブリック API、pandas DataFrame を受け取って返すという要件は、 ベクトル化された UDFs と同じです。ベクトル化された UDFs と同様に、これらの推論 APIs はPythonからSnowpark DataFrame を渡して呼び出すことができます。

public APIを持つカスタムモデルクラスのスケルトンを以下に示します。

context.path を使用してバイアスファイルを開き、 self.context.model_ref を使用して下位モデルクラスへの参照を取得すると、それらの predict メソッドを呼び出せるようになります。

from snowflake.ml.model import custom_model
import pandas as pd

class ExamplePipelineModel(custom_model.CustomModel):

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        ...
        return pd.DataFrame(...)
Copy

すべてのピースを組み合わせると、完全に機能するカスタムモデルとなります。

class ExamplePipelineModel(custom_model.CustomModel):
    def __init__(self, context: custom_model.ModelContext) -> None:
        super().__init__(context)
        v = open(context.path('config')).read()
        self.bias = json.loads(v)['bias']

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        features = self.context.model_ref('feature_preproc').transform(input)
        model_output = self.context.model_ref('m2').predict(
            self.context.model_ref('m1').predict(features)
        )
        return pd.DataFrame({
            'output': model_output + self.bias})
Copy

カスタムモデルの使用

新しいカスタムモデル(パイプライン)を以下のようにローカルで実行してテストすることができます。

my_model_pipeline = ExamplePipelineModel(mc)
output_df = my_model_pipeline.predict(input_df)
Copy

あるいは、レジストリにログを取り、Snowflakeに展開します。ここに示すように、モデルクラスが必要とするライブラリを指定するために conda_dependencies (もしくは pip_requirements)をプロバイダーとして指定します。

モデルの入力シグネチャを推論するために、 sample_input_data (1つのpandas DataFrame) を提供します。あるいは、 モデル署名 を提供します。

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model_pipeline,
            model_name="my_custom_model_pipeline",
            version_name="v1",
            conda_dependencies=["scikit-learn"],
            comment="My Custom ML Model Pipeline",
            sample_input_data=train_features)
output_df = mv.run(input_df)
Copy

テーブル機能推論

Snowpark ML 1.5.4では、複数の列を返す推論メソッドでモデルをログを記録できます。そのためには、 {"function_type": "TABLE_FUNCTION"} オプションをつけてモデルをログに記録し、上記のように @inference_api デコレーターを使用します。以下の例では、decorated メソッドは、2 つの出力列を含む pandas DataFrame を返します。

class ExampleTableFunctionModel(custom_model.CustomModel):

  @custom_model.inference_api
  def predict(self, input: pd.DataFrame) -> pd.DataFrame:
      output_df = pandas.DataFrame([{"OUTPUT1": input["INPUT1"] + 1, input["INPUT2"] + 1}])
      return output_df

my_model = ExampleTableFunctionModel()
reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model,
            model_name="my_custom_table_function_model",
            version_name="v1",
            options={"function_type": "TABLE_FUNCTION"},
            sample_input_data=train_features
            )
output_df = mv.run(input_df)
Copy

モデルに複数の推論手法が含まれている場合は、 method_options オプションを使用して、どれが FUNCTION でどれが TABLE_FUNCTION であるかを示しながらログに記録します。

reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
mv = reg.log_model(my_model,
            model_name="my_custom_table_function_model",
            version_name="v1",
            options={
              "method_options": {                                 ###
                "METHOD1": {"function_type": "TABLE_FUNCTION"},   ###
                "METHOD2": {"function_type": "FUNCTION"}          ###
              }
            },
            sample_input_data=train_features
            )
Copy

ログに記録されたモデルテーブル関数推論メソッドは、 SQL を介して以下のように呼び出すこともできます。

SELECT OUTPUT1, OUTPUT2
  FROM input_table,
      table(
          MY_MODEL!PREDICT(input_table.INPUT1, input_table.INPUT2)
      )
Copy

推論メソッドがデータパーティションを入力とし、複数の行と列を出力するパーティション・モデル推論メソッドの詳細については、 Partitioned Custom Models をご参照ください。