在 Snowflake Model Registry 中存储自定义模型¶

Snowflake Model Registry 允许您注册（记录）模型并在 Snowflake 内部使用它们进行推理。注册表支持多种类型的模型：

Snowpark ML Modeling
scikit-learn
XGBoost
LightGBM
CatBoost
PyTorch
TensorFlow
MLFlow PyFunc
Sentence Transformer
Hugging Face pipeline

模型注册表 API 还允许您记录其他类型的模型，包括使用外部工具训练的模型或从开源存储库获取的模型，只要它们是可序列化的，并且是从 snowflake.ml.model.custom_model.CustomModel 类派生的。

备注

本指南 (https://quickstarts.snowflake.com/guide/deploying_custom_models_to_snowflake_model_registry/index.html#0) 提供了记录自定义模型的示例。

此外，请查看这篇博文 (link removed) 以及随附的 Quickstarts，获取有关从 AWS SageMaker 或从 Azure ML 中导入模型的信息。

本主题介绍如何创建模型，将模型记录到 Snowflake Model Registry 中，以及如何部署模型。此功能的一个常见用例是定义多个类的管道，例如几个转换器或填补器后面跟着一个预测器或分类器。在这种情况下，自定义模型类本身相对简单，按顺序调用这些类，并将一个类的结果作为输入传递给下一个类。

定义模型上下文¶

除了代码之外，模型通常还需要一个或多个静态文件，例如配置文件或模型权重。自定义模型必须提供有关所有此类构件的信息，以便注册表知道将这些构件与模型一起打包。可以使用 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') 作为 Python 文件对象来读取工件（值 'config' 与传递给 ModelContext 的字典中的键相同）。

除了静态文件之外，模型还可以组成其他支持类型的模型或管道（例如，snowflake.ml.modeling.pipeline 或 scikit-learn 模型）。注册表已经知道如何记录这些类型的对象，因此您可以使用 model_refs 在模型上下文中直接传递 Python 对象，如下所示。您无需自己打包这些对象。这对于引导聚合算法（袋装算法）或预处理或后处理非常有用。

备注

model1 和 model2 是注册表本机支持的任何类型模型的对象。feature_preproc 是一个 snowflake.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') 等来检索这些从属模型的引用。如果模型公开了 predict 方法，您的代码可以直接从检索到的模型引用中调用该模型，例如使用 self.context.model_ref('m1').predict()。

总而言之，自定义模型的上下文声明要序列化的 Python 对象以及模型的工件，这些工件是模型使用的本地文件，必须与代码一起存储在 Snowflake 中。模型类使用上下文来查找代码和工件；无论您的模型是在本地运行还是在 Snowflake 中运行，这都适用。

编写自定义模型类¶

若要告诉模型注册表如何记录和部署自定义模型，请从 snowflake.ml.model.custom_model.CustomModel 继承。

模型可以公开多种推理方法（例如，scikit-learn 模型公开 predict 和 predict_proba 方法）。要在自定义模型中声明推理函数，请将它们定义为子类的公共方法，并用 @custom_model.inference_api 对其进行装饰。此装饰器指示方法是模型的公共 API 的一部分，允许通过模型注册表从 Python 或 SQL 调用它。使用 inference_api 装饰的方法必须接受并返回 Pandas DataFrame。可以使用 inference_api 装饰任意数量的方法。

备注

公共 API 接受和返回 Pandas DataFrame 的要求与矢量化 UDFs 的要求相同。与矢量化 UDFs 一样，这些推理 APIs 可以通过 Python 传递 Snowpark DataFrame 来调用

具有公共 API 的框架自定义模型类如下所示。

注意，使用 context.path 打开 bias 文件，然后使用 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

向 ModelContext 提供关键字实参¶

Snowflake ML 支持给 ModelContext 类传递任意数量的关键字实参，让开发者在定义和初始化自定义模型时轻松包含其他参数。

以下示例演示了如何向模型上下文提供关键字实参并在自定义模型类中使用它们：

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

# Initialize ModelContext with variable keyword arguments
mc = custom_model.ModelContext(
    config='local_model_dir/config.json',
    m1=model1
)

# Define a custom model class that utilizes the context
class ExamplePipelineModel(custom_model.CustomModel):
    def __init__(self, context: custom_model.ModelContext) -> None:
        super().__init__(context)
        # Access the 'config' file from the context
        with open(self.context['config'], 'r') as f:
            config_data = f.read()
        self.bias = json.loads(config_data)['bias']

    @custom_model.inference_api
    def predict(self, input: pd.DataFrame) -> pd.DataFrame:
        # Use the model 'm1' from the context to make predictions
        model_output = self.context['m1'].predict(input)
        # Apply bias from the context to the model output
        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 （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 装饰器。在以下示例中，装饰方法返回 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

要了解有关分区模型推理方法的更多信息，即推理方法将数据分区作为输入并输出多行和多列，请参阅分区自定义模型。