PyTorch

Die Snowflake ML Model Registry unterstützt Modelle, die mit PyTorch erstellt wurden (von torch.nn.Module abgeleitete Modelle).

Die folgenden zusätzlichen Optionen können im options-Dictionary verwendet werden, wenn Sie log_model abrufen:

Option

Beschreibung

target_methods

Liste mit den Namen der für das Modellobjekt verfügbaren Methoden. PyTorch-Modelle haben standardmäßig die folgende Zielmethode: forward.

cuda_version

Die Version der CUDA-Laufzeitumgebung, die beim Bereitstellen auf einer Plattform mit GPU verwendet werden soll. Der Standardwert ist 11.8. Wird das Modell manuell auf None gesetzt, kann es nicht auf einer Plattform mit GPU eingesetzt werden.

multiple_inputs

Ob das Modell mehrere Tensor-Eingaben erwartet. Die Standardeinstellung ist False. Wenn True, akzeptiert das Modell eine Liste von Tensoren als Eingabe anstelle eines einzelnen Tensors.

Sie müssen entweder den Parameter sample_input_data oder signatures angeben, wenn Sie ein PyTorch-Modell protokollieren, damit die Registry die Signaturen der Zielmethoden kennt.

Bemerkung

Bei Verwendung von pandas DataFrames (die standardmäßig float64 verwenden) stellen Sie sicher, dass Ihre PyTorch-Modellebenen mit dtype=torch.float64 erstellt werden, um dtype-Konfliktfehler zu vermeiden.

Beispiel

In diesem Beispiel wird davon ausgegangen, dass reg eine Instanz von snowflake.ml.registry.Registry ist.

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