PyTorch¶

O Snowflake ML Model Registry oferece suporte a modelos criados usando PyTorch (modelos derivados do torch.nn.Module).

As seguintes opções adicionais podem ser usadas no dicionário options quando você chamar log_model:


Opção	Descrição
`target_methods`	Uma lista dos nomes dos métodos disponíveis no objeto de modelo. Os modelos do PyTorch têm o seguinte método de destino por padrão: `forward`.
`cuda_version`	A versão do tempo de execução CUDA a ser usada ao implantar em uma plataforma com GPU; o padrão é 11.8. Se definido manualmente como `None`, o modelo não poderá ser implementado em uma plataforma com GPU.
`multiple_inputs`	Se o modelo espera ou não várias entradas de tensores. O padrão é `False`. Quando `True`, o modelo aceitará uma lista de tensores como entrada em vez de um único tensor.

Você deve especificar o parâmetro sample_input_data ou signatures ao registrar um modelo PyTorch para que o registro conheça as assinaturas dos métodos de destino.

Nota

Ao usar DataFrames pandas (que usam float64 por padrão), certifique-se de que as camadas do modelo do PyTorch sejam criadas com dtype=torch.float64 para evitar erros de incompatibilidade com dtype.

Exemplo¶

Este exemplo considera que reg é uma instância 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:])