PyTorch¶

Snowflake ML Model Registry 支持使用 PyTorch 创建的模型（从 torch.nn.Module 派生的模型）。

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

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

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

备注

使用 Pandas DataFrames（默认使用 float64）时，确保 PyTorch 模型层的创建方式为 dtype=torch.float64 以避免 dtype 不匹配错误。

示例¶

此示例假设 reg 是一个实例 snowflake.ml.registry.Registry。

import torch
import torch.nn as nn
from sklearn import datasets, model_selection

# Define a simple neural network for classification
class IrisClassifier(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        # Use float64 to match pandas DataFrame default dtype
        self.model = nn.Sequential(
            nn.Linear(input_dim, hidden_dim, dtype=torch.float64),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim, dtype=torch.float64),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim, dtype=torch.float64),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)

# 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)

# Create model
model = IrisClassifier(input_dim=4, hidden_dim=32, output_dim=3)

# Train the model
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

X_train_tensor = torch.tensor(X_train.values)
y_train_tensor = torch.tensor(y_train.values, dtype=torch.long)

model.train()
for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(X_train_tensor)
    loss = criterion(outputs, y_train_tensor)
    loss.backward()
    optimizer.step()

# 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:])

Copy