PyTorch

Snowflake MLモデルレジストリは、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_datasignatures のどちらかのパラメーターを指定する必要があります。

注釈

Pandas DataFramesを使用する場合(デフォルトではfloat64を使用)、dtype不一致エラーを回避するためにPyTorchモデルレイヤーが dtype=torch.float64 で作成されていることを確認してください。

この例では、 regsnowflake.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