PyTorch¶

Le registre de modèles Snowflake ML prend en charge les modèles créés à l’aide de PyTorch (modèles dérivés de torch.nn.Module).

Les options supplémentaires suivantes peuvent être utilisées dans le dictionnaire options lors de l’appel à log_model :


Option	Description
`target_methods`	Une liste des noms des méthodes disponibles sur l’objet modèle. Les modèles PyTorch ont la méthode cible suivante par défaut : `forward`.
`cuda_version`	La version de l’environnement d’exécution CUDA à utiliser lors du déploiement sur une plateforme avec GPU ; la valeur par défaut est 11.8. S’il est défini manuellement sur `None`, le modèle ne peut pas être déployé sur une plateforme ayant un GPU.
`multiple_inputs`	Si le modèle attend plusieurs entrées de tenseur. La valeur par défaut est `False`. Lorsque la valeur est `True`, le modèle accepte une liste de tenseurs en entrée au lieu d’un seul tenseur.

Vous devez spécifier le paramètre sample_input_data ou signatures lors de la journalisation d’un modèle PyTorch afin que le registre connaisse les signatures des méthodes cibles.

Note

Lorsque vous utilisez des pandas DataFrames (qui utilisent float64 par défaut), assurez-vous que les couches de votre modèle PyTorch sont créées avec dtype=torch.float64 pour éviter les erreurs de non correspondance dtype.

Exemple¶

Cet exemple suppose que reg est une instance de 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:])