Snowflake Model Registry¶
Note
The model registry API described in this topic is generally available as of package version 1.5.0.
The Snowflake Model Registry lets you securely manage models and their metadata in Snowflake, regardless of origin. The model registry stores machine learning models as first-class schema-level objects in Snowflake so they can easily be found and used by others in your organization. You can create registries, and store models in them, using Python classes in the Snowpark ML library. Models can have multiple versions, and you can designate a version as the default.
Once you have stored a model, you can invoke its methods (equivalent to functions or stored procedures) to perform model operations, such as inference, in a Snowflake virtual warehouse.
Tip
See Introduction to Machine Learning with Snowpark ML for an example of an end-to-end workflow in Snowpark ML, including the Snowflake Model Registry.
If you have models in Microsoft Azure Machine Learning or in Amazon SageMaker, see Deploying Models from Azure ML and SageMaker to Snowpark ML.
The most important classes in the Snowflake Model Registry Python API are:
snowflake.ml.registry.Registry: Manages models within a schema.
snowflake.ml.model.Model: Represents a model.
snowflake.ml.model.ModelVersion: Represents a version of a model.
The Snowflake Model Registry supports the following types of models.
scikit-learn
XGBoost
LightGBM
CatBoost
PyTorch
TensorFlow
MLFlow PyFunc
Sentence Transformer
Hugging Face pipeline
Other types of models via the
snowflake.ml.model.CustomModel
class (see Storing Custom Models in the Snowflake Model Registry).
This topic describes how to perform registry operations in Python using Snowpark ML. You can also perform many registry operations in SQL; see Model commands.
Differences from Private Preview¶
Snowflake previously made a model registry available privately to select customers. The registry feature described in this topic has significant changes to functionality and APIs compared to the private preview version. Most notably, the model registry functionality is now hosted natively inside Snowflake using a new schema-level object.
For details on the differences between these two APIs, see Migrating from the Snowflake Model Registry Preview API.
Required Privileges¶
To create a model, you must either own the schema where the model is created or have the CREATE MODEL privilege on it. To use a model, you must either own the model or have the USAGE privilege on it. The USAGE privilege allows grantees to use the model for inference without being able to see any of its internals.
If a user’s role has USAGE on a model, it appears in Snowsight’s model registry page. For details, see Access control privileges.
Note
Models currently do not support replication.
Current Limitations and Issues¶
The Snowflake Model Registry currently has the following limitations:
The registry cannot be used in Snowflake Native Apps.
Models cannot be shared or cloned, and are skipped during replication.
Versions 1.5.0 and 1.5.1 of the snowflake-ml-python
package have the following known issues. Until these are
addressed, use the provided workaround.
In Snowflake release 8.23 and earlier, the library does not work in owner’s rights stored procedures. Use caller’s rights stored procedures instead.
In stored procedures, logging a model requires embedding a copy of the local Snowpark ML library in the model. Specify the
embed_local_ml_library
option in thelog_model
call as follows.registry.log_model(..., options={"embed_local_ml_library": True, ...})
The following limits apply to models and model versions.
Models |
|
---|---|
Model versions |
|
Opening the Snowflake Model Registry¶
Models are first-class Snowflake objects and can be organized within a database and schema along with other Snowflake objects. The Snowflake Model Registry provides a Python class for managing models within a schema. Thus, any Snowflake schema can be used as a registry. It is not necessary to initialize or otherwise prepare a schema for this purpose. Snowflake recommends creating one or more dedicated schemas for this purpose, such as ML.REGISTRY. You can create the schema using CREATE SCHEMA.
Before you can create or modify models in the registry, you must open the registry. Opening the registry returns a reference to it, which you can then use to add new models and obtain references to existing models.
from snowflake.ml.registry import Registry
reg = Registry(session=sp_session, database_name="ML", schema_name="REGISTRY")
Registering Models and Versions¶
Adding a model to the registry is called logging the model. Log a model by calling the registry’s log_model
method. This method:
Serializes the model, a Python object, and creates a Snowflake model object from it.
Adds metadata, such as a description, to the model as specified in the
log_model
call.
Each model may have any number of versions. To log additional versions of the model, call log_model
again with
the same model_name
but a different version_name
.
You cannot add tags to a model when it is added to the registry, because tags are attributes of the model, and
log_model
adds a specific model version, only creating a model when adding its first version. You can update
the model’s tags after logging the first version of the model.
In the example below, clf
, short for “classifier,” is the Python model object, which has already been
created elsewhere in your code. You can add a comment at registration time, as shown here. The combination of
name and version must be unique in the schema. You may specify conda_dependencies
lists; the
specified packages will be deployed with the model.
mv = reg.log_model(clf,
model_name="my_model",
version_name="v1",
conda_dependencies=["scikit-learn"],
comment="My awesome ML model",
metrics={"score": 96},
sample_input_data=train_features)
The arguments of log_model
are described here.
Required Arguments
Argument |
Description |
---|---|
|
The Python model object of a supported model type. Must be serializable (“pickleable”). |
|
The model’s name, used with |
Note
The combination of model name and version must be unique in the schema.
Optional Arguments
Argument |
Description |
---|---|
|
String specifying the model’s version, used with |
|
List of paths to directories of code to import when loading or deploying the model. |
|
Comment, for example a description of the model. |
|
List of Conda packages required by your model. This argument specifies package names
and optional versions in Conda format,
that is, |
|
List of external modules to pickle with the model. Supported with scikit-learn, Snowpark ML, PyTorch, TorchScript, and custom models. |
|
Dictionary containing metrics linked to the model version. |
|
Dictionary containing options for model creation. The following options are available for all model types.
Individual model types may support additional options. See Notes on Specific Model Types. |
|
List of package specs for PyPI packages required by your model. |
|
The version of Python in which the model will run. Defaults to |
|
A DataFrame containing sample input data. The feature names required by the model, and their types, are
extracted from this DataFrame. Either this argument or |
|
Model method signatures as a mapping from target method name to signatures of input and output. Either this argument or
|
log_model
returns a snowflake.ml.model.ModelVersion
object, which represents the version of the model
that was added to the registry.
Once registered, the model itself cannot be modified (although you can change its metadata). To delete a model and all its versions, use the registry’s delete_model method.
Working with Model Artifacts¶
After a model has been logged, its artifacts (the files backing the model, including its serialized Python objects and various metadata files such as its manifest) are available on an internal stage. Artifacts cannot be modified, but you can view or download the artifacts of models you own.
Note
Having the USAGE privilege on a model does not allow you to access its artifacts; ownership is required.
You can access model artifacts from a stage using, for example, the GET command or its equivalent in Snowpark Python, FileOperation.get.
However, you cannot address model artifacts using the usual stage path syntax. Instead, use a snow://
URL, a more
general way to specify the location of objects in Snowflake. For example, a version inside a model can be specified by a
URL of the form snow://model/<model_name>/versions/<version_name>/
.
Knowing the of name of the model and the version you want, you can use the LIST command to view the artifacts of the model as follows:
LIST 'snow://model/my_model/versions/V3/';
The output resembles:
name size md5 last_modified
versions/V3/MANIFEST.yml 30639 2f6186fb8f7d06e737a4dfcdab8b1350 Thu, 18 Jan 2024 09:24:37 GMT
versions/V3/functions/apply.py 2249 e9df6db11894026ee137589a9b92c95d Thu, 18 Jan 2024 09:24:37 GMT
versions/V3/functions/predict.py 2251 132699b4be39cc0863c6575b18127f26 Thu, 18 Jan 2024 09:24:37 GMT
versions/V3/model.zip 721663 e92814d653cecf576f97befd6836a3c6 Thu, 18 Jan 2024 09:24:37 GMT
versions/V3/model/env/conda.yml 332 1574be90b7673a8439711471d58ec746 Thu, 18 Jan 2024 09:24:37 GMT
versions/V3/model/model.yaml 25718 33e3d9007f749bb2e98f19af2a57a80b Thu, 18 Jan 2024 09:24:37 GMT
To retrieve one of these artifacts, use the SQL GET command.
GET 'snow://model_my_model/versions/V3/MANIFEST.yml'
Or the equivalent with Snowpark Python:
session.file.get('snow://model_my_model/versions/V3/MANIFEST.yml', 'model_artifacts')
Note
The names and organization of a model’s artifacts can vary depending on the type of the model and might change. The example artifact list above is intended to be illustrative, not authoritative.
Deleting Models¶
Use the registry’s delete_model
method to delete a model and all its versions.
reg.delete_model("mymodel")
Tip
You can also delete models in SQL using DROP MODEL.
Getting Models from the Registry¶
To get information about each model, use the show_models
method.
model_df = reg.show_models()
Tip
In SQL, use SHOW MODELS to get a list of models.
The result of show_models
is a pandas DataFrame. The available columns are shown below.
Column |
Description |
---|---|
|
Date and time when the model was created. |
|
Name of the model. |
|
Database in which the model is stored. |
|
Schema in which the model is stored. |
|
Role that owns the model. |
|
Comment for the model. |
|
JSON array listing versions of the model. |
|
Version of the model used when referring to the model without a version. |
To get a list of the models in the registry instead, each as a Model
instance, use the models
method.
model_list = reg.models()
To get a reference to a specific model from the registry by name, use the registry’s get_model
method.
m = reg.get_model("MyModel")
Note
Model
instances are not copies of the original logged Python model object, but references to the underlying
model object in the registry.
Once you have a reference to a model, either one from the list returned by models
method or one retrieved using
get_model
, you can work with its metadata and
its versions.
Viewing and Updating a Model’s Metadata¶
You can view and update a model’s metadata attributes in the registry, including its name, comment, tags, and metrics.
Retrieving and Updating Comments¶
Use the model’s comment
attribute to retrieve and update the model’s comment.
print(m.comment)
m.comment = "A better description than the one I provided originally"
Note
The description
attribute is a synonym for comment
. The above code can also be written:
print(m.description)
m.description = "A better description than the one I provided originally"
Tip
You can also set a model’s comment in SQL using ALTER MODEL.
Renaming a Model¶
Use the rename
method to rename or move a model. Specify a fully qualified name as the new name to move the model to
a different database or schema.
m.rename("MY_MODEL_TOO")
Tip
You can also rename a model in SQL using ALTER MODEL.
Working with Model Versions¶
A model may have any number of versions, each identified by a string. You can use any version naming convention that you
like. Logging a model actually logs a specific version of the model. To log additional versions of a model, call
log_model
again with the same model_name
but a different version_name
.
Tip
In SQL, use SHOW VERSIONS IN MODEL to see the versions of a model.
A version of a model is represented by an instance of the snowflake.ml.model.ModelVersion
class.
To get a list of all the versions of a model, call the model object’s versions
method. The result is a list of
ModelVersion
instances.
version_list = m.versions()
To get information about each model as a DataFrame instead, call the model’s show_versions
method.
version_df = m.show_versions()
The resulting DataFrame contains the following columns.
Column |
Description |
---|---|
|
Date and time when the model version was created. |
|
Name of the version. |
|
Database in which the version is stored. |
|
Schema in which the version is stored. |
|
Name of the model that this version belongs to. |
|
Boolean value indicating whether this version is the model’s default version. |
|
JSON array of the names of the functions available in this version. |
|
JSON object containing metadata as key-value pairs ( |
|
JSON object from the |
Deleting Model Versions¶
You can delete a model version using the model’s delete_version
method.
m.delete_version("rc1")
Tip
You can also delete a model version in SQL using ALTER MODEL … DROP VERSION.
Default Version¶
A version of a model may be designated as the default model. Retrieve or set the model’s default
attribute to obtain
the current default version (as a ModelVersion
object) or to change it (using a string).
default_version = m.default
m.default = "v2"
Tip
In SQL, use ALTER MODEL to set the default version.
Model Version Aliases¶
You can assign an alias to a model version using the SQL ALTER MODEL command. You can use an alias wherever a version name is required, such as when getting a reference to a model version, in Python or in SQL.
In addition to aliases you create, the following system aliases are available in all models.
DEFAULT
refers to the default version of the model.FIRST
refers to the oldest version of the model by creation time.LAST
refers to the newest version of the model by creation time.
Alias names you create may not be the same as any existing version name or alias in the model, including system aliases.
Getting a Reference to a Model Version¶
To get a reference to a specific version of a model as a ModelVersion
instance, use the model’s version
method.
Use the model’s default
attribute to get the default version of the model.
mv = m.version("v1")
mv = m.default
Once you have a reference to a specific version of a model (such as the variable mv
in this example), you can
retrieve or update its comments or metrics and call the model’s methods (or functions) as shown in the following sections.
Retrieving and Updating Comments¶
As with models, model versions can have comments, which can be accessed and set via the model version’s comment
or
description
attribute.
print(mv.comment)
print(mv.description)
mv.comment = "A model version comment"
mv.description = "Same as setting the comment"
Tip
You can also change a model version’s comment in SQL using ALTER MODEL … MODIFY VERSION.
Retrieving and Updating Metrics¶
Metrics are key-value pairs used to track prediction accuracy and other model version characteristics. You can set
metrics when creating a model version or set them using the set_metric
method. A metric value can be any Python
object that can be serialized to JSON, including numbers, strings, lists, and dictionaries. Unlike tags, metric names
and possible values do not need to be defined in advance.
A test accuracy metric might be generated using sklearn’s accuracy_score
:
from sklearn import metrics
test_accuracy = metrics.accuracy_score(test_labels, prediction)
The confusion matrix can be generated similarly using sklearn:
test_confusion_matrix = metrics.confusion_matrix(test_labels, prediction)
Then we can set these values as metrics as follows.
# scalar metric
mv.set_metric("test_accuracy", test_accuracy)
# hierarchical (dictionary) metric
mv.set_metric("evaluation_info", {"dataset_used": "my_dataset", "accuracy": test_accuracy, "f1_score": f1_score})
# multivalent (matrix) metric
mv.set_metric("confusion_matrix", test_confusion_matrix)
To retrieve a model version’s metrics as a Python dictionary, use show_metrics
.
metrics = mv.show_metrics()
To delete a metric, call delete_metric
.
mv.delete_metric("test_accuracy")
Tip
You can also modify a model version’s metrics (which are stored in as metadata) in SQL using ALTER MODEL … MODIFY VERSION.
Retrieving Model Explanations¶
The model registry can explain a model’s results, telling you which input features contribute most to predictions, by calculating
Shapley values. This preview feature is available by default in all
model views created in Snowflake 8.31 and later through the underlying model’s explain
method. You can call explain
from SQL or via a model view’s
run
method in Python.
For details on this feature, see Model Explainability.
Exporting a Model Version¶
Use mv.export
to export a model’s files to a local directory. The directory is created if it does not exist.
mv.export("~/mymodel/")
By default, the exported files include the code, the environment to load the model, and model weights. To also
export the files needed to run the model in a warehouse, specify export_mode = ExportMode.FULL
.
mv.export("~/mymodel/", export_mode=ExportMode.FULL)
Loading a Model Version¶
Use mv.load
to load the original Python model object that was originally added to the registry. You can then
use the model for inference just as though you had defined it in your Python code.
clf = mv.load()
To ensure proper functionality of a model loaded from the registry, the target Python environment (that is, the
versions of the Python interpreter and of all libraries) should be identical to the environment from which the model
was logged. Specify force=True
in the load
call to force the model to be loaded even if the environment is
different.
Tip
To make sure your environment is the same as the one where the model is hosted, download a copy of the conda environment from the model registry:
conda_env = session.file.get("snow://model/<modelName>/versions/<versionName>/runtimes/python_runtime/env/conda.yml", ".")
open("~/conda.yml", "w").write(conda_env)
Then create a new conda environment from this file.
conda env create --name newenv --file=~/conda.yml
conda activate newenv
The optional options
argument is a dictionary of options for loading the model. Currently, the argument supports
only the use_gpu
option.
Option |
Type |
Description |
Default |
---|---|---|---|
|
|
Enables GPU-specific loading logic. |
|
The following example illustrates the use of the options
argument.
clf = mv.load(options={"use_gpu": True})
Calling Model Methods¶
Model versions can have methods, which are attached functions that can be executed to perform inference or other model operations. The versions of a model can have different methods, and the signatures of these methods can also differ.
To call a method of a model version, use mv.run
, specifying the name of the function to be called and passing a
Snowpark or pandas DataFrame containing the inference data and any other required parameters. The method executes in a
Snowflake warehouse.
The return value of the method is a Snowpark or pandas DataFrame, depending on the type of DataFrame passed in.
Snowpark DataFrames are evaluated lazily, so the method is run only when the DataFrame’s collect
, show
,
or to_pandas
method is called.
Note
Invoking a method runs it in the warehouse specified in the session you’re using to connect to the registry. See Specifying a Warehouse.
The following example illustrates running the predict
method of a model. This model’s predict
method does not
require any parameters besides the inference data (test_features
here). If it did, they would be passed as
additional arguments after the inference data.
remote_prediction = mv.run(test_features, function_name="predict")
remote_prediction.show() # assuming test_features is Snowpark DataFrame
To see what methods can be called on a given model, call mv.show_functions
. The return value of this method is a
list of ModelFunctionInfo
objects. Each of these objects includes the following attributes:
name
: The name of the function that can be called from Python or SQL.target_method
: The name of the Python method in the original logged model.
Tip
You can also call model methods in SQL. See Model methods.
Cost Considerations¶
Using the Snowflake Model Registry incurs standard Snowflake consumption-based costs. These include:
Cost of storing model artifacts, metadata, and functions. See Exploring storage cost for general information about storage costs.
Cost of copying files between stages to Snowflake. See COPY FILES.
Cost of serverless model object operations through the Snowsight UI or the SQL or Python interface, such as showing models and model versions, and altering model comments, tags, and metrics.
Warehouse compute costs, which vary depending on the type of model and the quantity of data used in inference. See Understanding compute cost for general information about Snowflake compute costs. Warehouse compute costs are incurred for:
Model and version creation operations.
Invoking a model’s methods.
Notes on Specific Model Types¶
This section provides additional information on logging specific types of models into the Snowflake Model Registry.
Snowpark ML¶
The registry supports models created using Snowpark ML modeling APIs (models derived from
snowpark.ml.modeling.framework.base.BaseEstimator
). The following additional options may be used in the options
dictionary
when calling log_model
.
Option |
Description |
---|---|
|
A list of the names of the methods available on the model object. Snowpark ML models have the following target methods by default, assuming the method exists:
|
You do not need to specify sample_input_data
or signatures
when logging a Snowpark ML model;
these are automatically inferred during fitting.
Example¶
import pandas as pd
import numpy as np
from sklearn import datasets
from snowflake.ml.modeling.xgboost import XGBClassifier
iris = datasets.load_iris()
df = pd.DataFrame(data=np.c_[iris["data"], iris["target"]], columns=iris["feature_names"] + ["target"])
df.columns = [s.replace(" (CM)", "").replace(" ", "") for s in df.columns.str.upper()]
input_cols = ["SEPALLENGTH", "SEPALWIDTH", "PETALLENGTH", "PETALWIDTH"]
label_cols = "TARGET"
output_cols = "PREDICTED_TARGET"
clf_xgb = XGBClassifier(
input_cols=input_cols, output_cols=output_cols, label_cols=label_cols, drop_input_cols=True
)
clf_xgb.fit(df)
model_ref = registry.log_model(
clf_xgb,
model_name="XGBClassifier",
version_name="v1",
)
model_ref.run(df.drop(columns=label_cols).head(10), function_name='predict_proba')
scikit-learn¶
The registry supports models created using scikit-learn (models derived from sklearn.base.BaseEstimator
or
sklearn.pipeline.Pipeline
). The following additional options may be used in the options
dictionary
when calling log_model
.
Option |
Description |
---|---|
|
A list of the names of the methods available on the model object. scikit-learn models have the following target methods by default, assuming the method exists:
|
You must specify either the sample_input_data
or signatures
parameter when logging a scikit-learn model
so that the registry knows the signatures of the target methods.
Example¶
from sklearn import datasets, ensemble
iris_X, iris_y = datasets.load_iris(return_X_y=True, as_frame=True)
clf = ensemble.RandomForestClassifier(random_state=42)
clf.fit(iris_X, iris_y)
model_ref = registry.log_model(
clf,
model_name="RandomForestClassifier",
version_name="v1",
sample_input_data=iris_X,
options={
"method_options": {
"predict": {"case_sensitive": True},
"predict_proba": {"case_sensitive": True},
"predict_log_proba": {"case_sensitive": True},
}
},
)
model_ref.run(iris_X[-10:], function_name='"predict_proba"')
XGBoost¶
The registry supports models created using XGBoost (models derived from xgboost.XGBModel
or xgboost.Booster
).
The following additional options may be used in the options
dictionary when calling log_model
.
Option |
Description |
---|---|
|
A list of the names of the methods available on the model object. Models derived from |
|
The version of the CUDA runtime to be used when deploying to a platform with GPU; defaults to 11.7. If manually set
to |
You must specify either the sample_input_data
or signatures
parameter when logging an XGBoost model so
that the registry knows the signatures of the target methods.
Example¶
import xgboost
from sklearn import datasets, model_selection
cal_X, cal_y = datasets.load_breast_cancer(as_frame=True, return_X_y=True)
cal_X_train, cal_X_test, cal_y_train, cal_y_test = model_selection.train_test_split(cal_X, cal_y)
params = dict(n_estimators=100, reg_lambda=1, gamma=0, max_depth=3, objective="binary:logistic")
regressor = xgboost.train(params, xgboost.DMatrix(data=cal_X_train, label=cal_y_train))
model_ref = registry.log_model(
regressor,
model_name="xgBooster",
version_name="v1",
sample_input_data=cal_X_test,
options={
"target_methods": ["predict"],
"method_options": {
"predict": {"case_sensitive": True},
},
},
)
model_ref.run(cal_X_test[-10:])
PyTorch¶
The registry supports PyTorch models (classes derived from torch.nn.Module
or torch.jit.ModuleScript
) if the
model’s forward
method accepts one or more torch.Tensor
instances as input and returns a torch.Tensor
or a
tuple of them. The registry converts between pandas DataFrames and tensors when calling the model and returning the results.
Tensors correspond to columns in the dataframe.
For example, suppose your model accepts two tensors like this:
import torch
class TorchModel(torch.nn.Module):
def __init__(self, n_input: int, n_hidden: int, n_out: int, dtype: torch.dtype = torch.float32) -> None:
super().__init__()
self.model = torch.nn.Sequential(
torch.nn.Linear(n_input, n_hidden, dtype=dtype),
torch.nn.ReLU(),
torch.nn.Linear(n_hidden, n_out, dtype=dtype),
torch.nn.Sigmoid(),
)
def forward(self, tensor_1: torch.Tensor, tensor_2: torch.Tensor) -> torch.Tensor:
return self.model(tensor_1) + self.model(tensor_2)
If you want to pass torch.Tensor([[1,2],[3,4]])
as tensor_1
and torch.Tensor([[5,6], [7,8]])
as
tensor_2
, create a DataFrame as follows to pass to the model.
import pandas as pd
tensors = pd.DataFrame([[[1,2],[5,6]],[[3,4],[7,8]]])
Then the tensors
DataFrame looks like this.
0 1
0 [1, 2] [5, 6]
1 [3, 4] [7, 8]
Similarly, if your model returns two tensors, such as (torch.Tensor([[1,2],[3,4]]), torch.Tensor([[5,6], [7,8]]))
,
the result is a DataFrame like the one above.
When providing sample input data for a PyTorch model, you must provide either a list of tensors (which will be converted ta a pandas DataFrame) or a DataFrame. A list may contain a single tensor, but a tensor on its own is not accepted.
Logging the model¶
The following additional options may be used in the options
dictionary when calling log_model
.
Option |
Description |
---|---|
|
A list of the names of the methods available on the model object. PyTorch models default to |
|
The version of the CUDA runtime to be used when deploying to a platform with GPU; defaults to 11.7. If manually set
to |
You must specify either the sample_input_data
or signatures
parameter when logging a PyTorch model
so that the registry knows the signatures of the target methods.
Example¶
import torch
import numpy as np
class TorchModel(torch.nn.Module):
def __init__(self, n_input: int, n_hidden: int, n_out: int, dtype: torch.dtype = torch.float32) -> None:
super().__init__()
self.model = torch.nn.Sequential(
torch.nn.Linear(n_input, n_hidden, dtype=dtype),
torch.nn.ReLU(),
torch.nn.Linear(n_hidden, n_out, dtype=dtype),
torch.nn.Sigmoid(),
)
def forward(self, tensor: torch.Tensor) -> torch.Tensor:
return self.model(tensor)
n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01
dtype = torch.float32
x = np.random.rand(batch_size, n_input)
data_x = torch.from_numpy(x).to(dtype=dtype)
data_y = (torch.rand(size=(batch_size, 1)) < 0.5).to(dtype=dtype)
model = TorchModel(n_input, n_hidden, n_out, dtype=dtype)
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
for _epoch in range(100):
pred_y = model.forward(data_x)
loss = loss_function(pred_y, data_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
model_ref = registry.log_model(
model,
model_name="torchModel",
version_name="v1",
sample_input_data=[data_x],
)
model_ref.run([data_x])
TensorFlow¶
Models that extend tensorflow.Module
or tensorflow.keras.Model
are supported when they accept and return tensors
and are compilable or compiled.
The
__call__
method for atensorflow.Module
or thecall
method for atensorflow.keras.Model
accepts one or moretensorflow.Tensor
ortensorflow.Variable
as input and returns atensorflow.Tensor
ortensorflow.Variable
or a tuple of one of these types.If your model extends
Module
, it must be compilable, meaning the__call__
method is decorated with@tensorflow.function
; see tf.function documentation. If it extendsModel
, it must be compiled; see compile documentation.
The registry converts between pandas DataFrames and tensors when calling the model and returning the results. Tensors correspond to columns in the dataframe.
For example, suppose your model accepts two tensors like this:
import tensorflow as tf
class KerasModel(tf.keras.Model):
def __init__(self, n_hidden: int, n_out: int) -> None:
super().__init__()
self.fc_1 = tf.keras.layers.Dense(n_hidden, activation="relu")
self.fc_2 = tf.keras.layers.Dense(n_out, activation="sigmoid")
def call(self, tensor_1: tf.Tensor, tensor_2: tf.Tensor) -> tf.Tensor:
input = tensor_1 + tensor_2
x = self.fc_1(input)
x = self.fc_2(x)
return x
If you want to pass tf.Tensor([[1,2],[3,4]])
as tensor_1
and tf.Tensor([[5,6], [7,8]])
as
tensor_2
, create a DataFrame as follows to pass to the model.
import pandas as pd
tensors = pd.DataFrame([[[1,2],[5,6]],[[3,4],[7,8]]])
Then the tensors
DataFrame looks like this.
0 1
0 [1, 2] [5, 6]
1 [3, 4] [7, 8]
Similarly, if your model returns two tensors, such as (tf.Tensor([[1,2],[3,4]]), tf.Tensor([[5,6], [7,8]]))
,
the result is a DataFrame like the one above.
When providing sample input data for a TensorFlow model, you must provide either a list of tensors (which will be converted ta a pandas DataFrame) or a DataFrame. A list may contain a single tensor, but a tensor on its own is not accepted.
Logging the model¶
The following additional options may be used in the options
dictionary when calling log_model
.
Option |
Description |
---|---|
|
A list of the names of the methods available on the model object. TensorFlow models default to |
|
The version of the CUDA runtime to be used when deploying to a platform with GPU; defaults to 11.7. If manually set
to |
You must specify either the sample_input_data
or signatures
parameter when logging a TensorFlow model
so that the registry knows the signatures of the target methods.
Example¶
import tensorflow as tf
import numpy as np
class KerasModel(tf.keras.Model):
def __init__(self, n_hidden: int, n_out: int) -> None:
super().__init__()
self.fc_1 = tf.keras.layers.Dense(n_hidden, activation="relu")
self.fc_2 = tf.keras.layers.Dense(n_out, activation="sigmoid")
def call(self, tensor: tf.Tensor) -> tf.Tensor:
input = tensor
x = self.fc_1(input)
x = self.fc_2(x)
return x
n_input, n_hidden, n_out, batch_size, learning_rate = 10, 15, 1, 100, 0.01
dtype = tf.float32
x = np.random.rand(batch_size, n_input)
data_x = tf.convert_to_tensor(x, dtype=dtype)
raw_data_y = tf.random.uniform((batch_size, 1))
raw_data_y = tf.where(raw_data_y > 0.5, tf.ones_like(raw_data_y), tf.zeros_like(raw_data_y))
data_y = tf.cast(raw_data_y, dtype=dtype)
model = KerasModel(n_hidden, n_out)
model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate), loss=tf.keras.losses.MeanSquaredError())
model.fit(data_x, data_y, batch_size=batch_size, epochs=100)
model_ref = registry.log_model(
model,
model_name="tfModel",
version_name="v1",
sample_input_data=[data_x],
)
model_ref.run([data_x])
MLFlow¶
MLFlow models that provide a PyFunc flavor are supported. If your MLFlow model has a signature, the signature
argument is inferred from the model. Otherwise, you must provide either signature
or sample_input_data
.
The following additional options may be used in the options
dictionary when calling log_model
.
Option |
Description |
---|---|
|
The URI of the artifacts of the MLFlow model. Must be provided if it is not available in the model’s metadata as
|
|
If |
|
If |
Example¶
import mlflow
from sklearn import datasets, model_selection, ensemble
db = datasets.load_diabetes(as_frame=True)
X_train, X_test, y_train, y_test = model_selection.train_test_split(db.data, db.target)
with mlflow.start_run() as run:
rf = ensemble.RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
rf.fit(X_train, y_train)
# Use the model to make predictions on the test dataset.
predictions = rf.predict(X_test)
signature = mlflow.models.signature.infer_signature(X_test, predictions)
mlflow.sklearn.log_model(
rf,
"model",
signature=signature,
)
run_id = run.info.run_id
model_ref = registry.log_model(
mlflow.pyfunc.load_model(f"runs:/{run_id}/model"),
model_name="mlflowModel",
version_name="v1",
conda_dependencies=["mlflow<=2.4.0", "scikit-learn", "scipy"],
options={"ignore_mlflow_dependencies": True}
)
model_ref.run(X_test)
Hugging Face Pipeline¶
Note
For details on the expected input and output of specific types of Hugging Face pipelines, see Inferred Signatures for Hugging Face Pipelines.
The registry supports Hugging Face model classes defined as transformers that derive from transformers.Pipeline
.
The following code is an example of logging a compatible model.
lm_hf_model = transformers.pipeline(
task="text-generation",
model="bigscience/bloom-560m",
token="...", # Put your HuggingFace token here.
return_full_text=False,
max_new_tokens=100,
)
lmv = reg.log_model(lm_hf_model, model_name='bloom', version_name='v560m')
Important
A model based on huggingface_pipeline.HuggingFacePipelineModel
contains only configuration data; the model
weights are downloaded from the Hugging Face Hub each time the model is used.
The model registry currently supports deploying models only to warehouses. Warehouses do not support external network access without special configuration. Even if the required external access integrations have been created, there is no way at this time to specify which integrations a particular model needs.
The current best practice is to instead use transformers.Pipeline
as shown in the example above. This downloads
model weights to your local system and uploads the entire model to the warehouse. This results in a self-contained
model that does not need Internet access.
The registry infers the signatures
argument as long as the pipeline contains only a task from the following list.
conversational
fill-mask
question-answering
summarization
table-question-answering
text2text-generation
text-classification
(also calledsentiment-analysis
)text-generation
token-classification
(also calledner
)translation
translation_xx_to_yy
zero-shot-classification
The sample_input_data
argument is completely ignored for Hugging Face models. Specify the signatures
argument
when logging a Hugging Face model that is not in the above list so that the registry knows the signatures of the target
methods.
To see the inferred signature, use the show_functions
method. The following, for example, is the result of
lmv.show_functions()
where lmv
is the model logged above.
{'name': '__CALL__',
'target_method': '__call__',
'signature': ModelSignature(
inputs=[
FeatureSpec(dtype=DataType.STRING, name='inputs')
],
outputs=[
FeatureSpec(dtype=DataType.STRING, name='outputs')
]
)}]
With this information, you can call the model as follows.
import pandas as pd
remote_prediction = lmv.run(pd.DataFrame(["Hello, how are you?"], columns=["inputs"]))
Usage Notes¶
Many Hugging Face models are large and do not fit in a standard warehouse. Use a Snowpark-optimized warehouse or choose a smaller version of the model. For example, instead of using the
Llama-2-70b-chat-hf
model, tryLlama-2-7b-chat-hf
.Snowflake warehouses do not have GPUs. Use only CPU-optimized Hugging Face models.
Some Hugging Face transformers return an array of dictionaries per input row. The registry converts such output to a string containing a JSON representation of the array. For example, multi-output Question Answering output looks like this:
[{"score": 0.61094731092453, "start": 139, "end": 178, "answer": "learn more about the world of athletics"}, {"score": 0.17750297486782074, "start": 139, "end": 180, "answer": "learn more about the world of athletics.\""}]
You must specify either the sample_input_data
or signatures
parameter when logging a Hugging Face model
so that the registry knows the signatures of the target methods.
Example¶
# Prepare model
import transformers
import pandas as pd
finbert_model = transformers.pipeline(
task="text-classification",
model="ProsusAI/finbert",
top_k=2,
)
# Log the model
mv = registry.log_model(
finbert_model,
model_name="finbert",
version_name="v1",
)
# Use the model
mv.run(pd.DataFrame(
[
["I have a problem with my Snowflake that needs to be resolved asap!!", ""],
["I would like to have udon for today's dinner.", ""],
]
)
)
Result:
0 [{"label": "negative", "score": 0.8106237053871155}, {"label": "neutral", "score": 0.16587384045124054}]
1 [{"label": "neutral", "score": 0.9263970851898193}, {"label": "positive", "score": 0.05286872014403343}]