TensorFlow¶

Snowflake ML Model Registry 支持使用 TensorFlow 创建的模型（从 tensorflow.Module 派生的模型）和 Keras v2 模型（keras.Model，Keras 版本低于 3.0.0）。

备注

对于 Keras 3.0.0 或更高版本，请使用 Keras 处理程序。

调用 options 时，可以在 log_model 字典中使用下列附加选项：


选项	描述
`target_methods`	模型对象上可用方法的名称列表。TensorFlow 模型将 `__call__` 作为默认目标方法。Keras v2 模型将 `predict` 作为默认目标方法。
`cuda_version`	部署到具有 GPU 的平台时使用的 CUDA 运行时版本；默认值为 11.8。如果手动设置为 `None`，则无法将模型部署到具有 GPU 的平台。
`multiple_inputs`	模型是否预期接收多个张量输入。默认为 `False`。为 `True` 时，模型将接受张量列表作为输入，而不是单个张量。

在登记 TensorFlow 模型时，您必须指定 sample_input_data 或 signatures 参数，以确保注册表了解目标方法的签名。

备注

Keras v2 模型只能有一种目标方法。

备注

使用 Pandas DataFrames（默认使用 float64）时，确保 TensorFlow 模型使用 tf.float64 变量和 tf.TensorSpec 输入签名以避免 dtype 不匹配错误。

示例¶

这些示例假设 reg 是 snowflake.ml.registry.Registry 的一个实例。

TensorFlow 模块¶

以下示例演示了如何通过子类化 tf.Module 创建 TensorFlow 模型、将其记录到 Snowflake ML Model Registry 并运行推理。

import tensorflow as tf
import pandas as pd

# Define a simple TensorFlow module
class LinearModel(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.weight = tf.Variable(2.0, dtype=tf.float64, name="weight")
        self.bias = tf.Variable(1.0, dtype=tf.float64, name="bias")

    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 1), dtype=tf.float64)])
    def __call__(self, x):
        return self.weight * x + self.bias

# Create model instance
model = LinearModel(name="linear_model")

# Create sample input data as DataFrame
sample_df = pd.DataFrame({"input": [1.0, 2.0, 3.0, 4.0, 5.0]})

# Log the model
model_ref = reg.log_model(
    model=model,
    model_name="my_tf_linear_model",
    version_name="v1",
    sample_input_data=sample_df,
)

# Make predictions (default target method is __call__)
test_df = pd.DataFrame({"input": [6.0, 7.0, 8.0]})
result_df = model_ref.run(test_df)

Keras v2 序列模型¶

以下示例演示了如何训练 Keras v2 序列模型、将其记录到 Snowflake ML Model Registry 并运行推理。

import tf_keras as keras
from sklearn import datasets, model_selection

# Load dataset
iris = datasets.load_iris(as_frame=True)
X = iris.data
y = iris.target

# Rename columns for valid Snowflake identifiers
X.columns = [col.replace(' ', '_').replace('(', '').replace(')', '') for col in X.columns]

X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=0.2)

# Build Keras v2 model
model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(X_train.shape[1],)),
    keras.layers.Dense(32, activation='relu'),
    keras.layers.Dense(3, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Train the model
model.fit(X_train, y_train, epochs=50, verbose=0)

# Log the model
model_ref = reg.log_model(
    model=model,
    model_name="my_iris_classifier",
    version_name="v1",
    sample_input_data=X_test,
)

# Make predictions
result_df = model_ref.run(X_test[-10:], function_name="predict")