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:])
Copy
Langage: Français