Ajusta un modelo CIFAR-10 con Ray Tune

Abrir la versión 'notebook' de esta página

Utilice Ray Tune y el algoritmo asíncrono de reducción sucesiva a la mitad (ASHA) para buscar hiperparámetros para un clasificador de imágenes de PyTorch en el recurso de proceso 1xA10 AI Runtime conectado. Este cuaderno muestra cómo:

  • Inicia Ray con el cómputo de GPU adjunto.
  • Definir una función de entrenamiento de CIFAR-10 con puntos de control.
  • Ejecuta cuatro ensayos de Ray Tune, con dos ensayos de GPU fraccional funcionando simultáneamente.
  • Selecciona y evalúa el mejor ensayo.

Note

Este ejemplo requiere el entorno de IA de Databricks versión 6 o superior.

Prerequisites

Conecta este cuaderno al cómputo con GPU de AI Runtime:

  1. Selecciona Conectar en la parte superior del cuaderno.
  2. Seleccione GPU sin servidor.
  3. En el panel lateral Entorno, configura Acelerador en 1xA10.
  4. Selecciona AI v6 como entorno base.
  5. Selecciona Aplicar y luego selecciona Confirmar.

AI v6 incluye Ray, PyTorch y torchvision, por lo que este ejemplo no instala paquetes adicionales. La primera partida descarga el conjunto de datos CIFAR-10.

Inicializar Ray

Inicia Ray para la sesión del cuaderno con ray_init(). La función imprime un enlace de panel que funciona a través del proxy del controlador Databricks.

import ray
from serverless_gpu import ray_init

ray_init()

cluster_resources = ray.cluster_resources()
if cluster_resources.get("GPU", 0) < 1:
    raise RuntimeError(
        "Ray did not detect a GPU. Attach the notebook to 1xA10 compute, then rerun it."
    )
cluster_resources

El ejemplo limita cada ensayo a un subconjunto de CIFAR-10 y cinco épocas de entrenamiento para reducir el tiempo de ejecución de la búsqueda. Cada prueba reserva la mitad de la GPU A10, lo que permite a Ray programar hasta dos ensayos simultáneamente.

import tempfile
import uuid
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from filelock import FileLock
from ray import tune
from ray.tune import Checkpoint
from ray.tune.schedulers import ASHAScheduler
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

SEED = 42
MAX_EPOCHS = 5
NUM_SAMPLES = 4
CPUS_PER_TRIAL = 2
GPUS_PER_TRIAL = 0.5
TRAIN_SAMPLE_COUNT = 10_000
VALIDATION_SAMPLE_COUNT = 2_000
TEST_SAMPLE_COUNT = 2_000

DATA_DIR = Path(tempfile.gettempdir()) / "ray-tune-cifar10-data"
RESULTS_DIR = Path(tempfile.gettempdir()) / "ray-tune-cifar10-results"

Prepara el conjunto de datos y el modelo

El bloqueo de archivo impide que las pruebas simultáneas descarguen CIFAR-10 en el mismo directorio al mismo tiempo. Cada ensayo utiliza los mismos subconjuntos deterministas de entrenamiento y validación.

def load_cifar10(data_dir: Path):
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )

    data_dir.mkdir(parents=True, exist_ok=True)
    with FileLock(str(data_dir / "download.lock")):
        full_train_dataset = datasets.CIFAR10(
            root=data_dir, train=True, download=True, transform=transform
        )
        full_test_dataset = datasets.CIFAR10(
            root=data_dir, train=False, download=True, transform=transform
        )

    split_generator = torch.Generator().manual_seed(SEED)
    remaining_train_count = (
        len(full_train_dataset) - TRAIN_SAMPLE_COUNT - VALIDATION_SAMPLE_COUNT
    )
    train_dataset, validation_dataset, _ = random_split(
        full_train_dataset,
        [TRAIN_SAMPLE_COUNT, VALIDATION_SAMPLE_COUNT, remaining_train_count],
        generator=split_generator,
    )
    test_dataset, _ = random_split(
        full_test_dataset,
        [TEST_SAMPLE_COUNT, len(full_test_dataset) - TEST_SAMPLE_COUNT],
        generator=torch.Generator().manual_seed(SEED),
    )
    return train_dataset, validation_dataset, test_dataset


class CifarNet(nn.Module):
    def __init__(self, first_hidden_size: int, second_hidden_size: int):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, first_hidden_size)
        self.fc2 = nn.Linear(first_hidden_size, second_hidden_size)
        self.fc3 = nn.Linear(second_hidden_size, 10)

    def forward(self, inputs):
        inputs = self.pool(F.relu(self.conv1(inputs)))
        inputs = self.pool(F.relu(self.conv2(inputs)))
        inputs = torch.flatten(inputs, 1)
        inputs = F.relu(self.fc1(inputs))
        inputs = F.relu(self.fc2(inputs))
        return self.fc3(inputs)

Defina la función de entrenamiento de Ray Tune

Ray Tune llama a esta función para cada configuración de hiperparámetros muestreada. Al final de cada época, la función informa de métricas de validación y un punto de control. ASHA utiliza la pérdida reportada para evitar que los ensayos tengan un rendimiento inferior a tiempo.

def train_cifar(config, data_dir: Path, max_epochs: int):
    torch.manual_seed(SEED)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = CifarNet(config["first_hidden_size"], config["second_hidden_size"])
    model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        model.parameters(), lr=config["learning_rate"], momentum=0.9
    )
    start_epoch = 0

    incoming_checkpoint = tune.get_checkpoint()
    if incoming_checkpoint:
        with incoming_checkpoint.as_directory() as checkpoint_dir:
            checkpoint_state = torch.load(
                Path(checkpoint_dir) / "checkpoint.pt",
                map_location=device,
                weights_only=True,
            )
        model.load_state_dict(checkpoint_state["model_state"])
        optimizer.load_state_dict(checkpoint_state["optimizer_state"])
        start_epoch = checkpoint_state["epoch"] + 1

    train_dataset, validation_dataset, _ = load_cifar10(data_dir)
    train_loader = DataLoader(
        train_dataset, batch_size=config["batch_size"], shuffle=True
    )
    validation_loader = DataLoader(
        validation_dataset, batch_size=256, shuffle=False
    )

    for epoch in range(start_epoch, max_epochs):
        model.train()
        for inputs, labels in train_loader:
            inputs = inputs.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)
            optimizer.zero_grad()
            loss = criterion(model(inputs), labels)
            loss.backward()
            optimizer.step()

        model.eval()
        validation_loss = 0.0
        correct_predictions = 0
        prediction_count = 0
        with torch.no_grad():
            for inputs, labels in validation_loader:
                inputs = inputs.to(device, non_blocking=True)
                labels = labels.to(device, non_blocking=True)
                outputs = model(inputs)
                validation_loss += criterion(outputs, labels).item() * labels.size(0)
                correct_predictions += (outputs.argmax(dim=1) == labels).sum().item()
                prediction_count += labels.size(0)

        metrics = {
            "loss": validation_loss / prediction_count,
            "accuracy": correct_predictions / prediction_count,
        }
        with tempfile.TemporaryDirectory() as checkpoint_dir:
            torch.save(
                {
                    "epoch": epoch,
                    "model_state": model.state_dict(),
                    "optimizer_state": optimizer.state_dict(),
                },
                Path(checkpoint_dir) / "checkpoint.pt",
            )
            tune.report(
                metrics, checkpoint=Checkpoint.from_directory(checkpoint_dir)
            )

La búsqueda explora los tamaños de las capas totalmente conectadas, la tasa de aprendizaje y el tamaño del lote. El planificador ASHA puede detener un ensayo tras su primera época cuando la pérdida de validación probablemente no competirá con ensayos mejores.

search_space = {
    "first_hidden_size": tune.choice([64, 128, 256]),
    "second_hidden_size": tune.choice([32, 64, 128]),
    "learning_rate": tune.loguniform(1e-4, 1e-1),
    "batch_size": tune.choice([64, 128, 256]),
}

asha_scheduler = ASHAScheduler(
    max_t=MAX_EPOCHS,
    grace_period=1,
    reduction_factor=2,
)

trainable = tune.with_resources(
    tune.with_parameters(
        train_cifar, data_dir=DATA_DIR, max_epochs=MAX_EPOCHS
    ),
    resources={"cpu": CPUS_PER_TRIAL, "gpu": GPUS_PER_TRIAL},
)

Antes de ejecutar la siguiente celda, abre el enlace del panel impreso por ray_init(). En la página Jobs, abre el trabajo en ejecución para inspeccionar los actores del ensayo, los registros de tareas y las reservas de GPU. Con una reserva de GPU de 0,5 por prueba, Ray puede ejecutar dos pruebas simultáneamente en computación 1xA10.

tuner = tune.Tuner(
    trainable,
    param_space=search_space,
    tune_config=tune.TuneConfig(
        metric="loss",
        mode="min",
        scheduler=asha_scheduler,
        num_samples=NUM_SAMPLES,
        max_concurrent_trials=2,
    ),
    run_config=tune.RunConfig(
        name=f"cifar10-{uuid.uuid4().hex[:8]}",
        storage_path=str(RESULTS_DIR),
        checkpoint_config=tune.CheckpointConfig(
            num_to_keep=1,
            checkpoint_score_attribute="loss",
            checkpoint_score_order="min",
        ),
    ),
)
results = tuner.fit()

if results.num_errors:
    trial_errors = "\n\n".join(str(error) for error in results.errors)
    raise RuntimeError(
        f"{results.num_errors} of {len(results)} trials failed:\n{trial_errors}"
    )

Examina la mejor prueba

Selecciona el ensayo con la menor pérdida de validación reportada e inspecciona sus hiperparámetros y métricas finales.

best_result = results.get_best_result(metric="loss", mode="min")

print("Best hyperparameters:", best_result.config)
print(f"Validation loss: {best_result.metrics['loss']:.4f}")
print(f"Validation accuracy: {best_result.metrics['accuracy']:.2%}")

results_dataframe = results.get_dataframe()
display(
    results_dataframe[
        [
            "loss",
            "accuracy",
            "config/first_hidden_size",
            "config/second_hidden_size",
            "config/learning_rate",
            "config/batch_size",
        ]
    ].sort_values("loss")
)

Evalúa el mejor punto de control

Carga el punto de control de la prueba seleccionada y evalúalo en un subconjunto de prueba reservado de CIFAR-10.

best_model = CifarNet(
    best_result.config["first_hidden_size"],
    best_result.config["second_hidden_size"],
)
evaluation_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
best_model.to(evaluation_device)

with best_result.checkpoint.as_directory() as checkpoint_dir:
    checkpoint_state = torch.load(
        Path(checkpoint_dir) / "checkpoint.pt",
        map_location=evaluation_device,
        weights_only=True,
    )
best_model.load_state_dict(checkpoint_state["model_state"])
best_model.eval()

_, _, test_dataset = load_cifar10(DATA_DIR)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)
correct_predictions = 0
prediction_count = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        inputs = inputs.to(evaluation_device, non_blocking=True)
        labels = labels.to(evaluation_device, non_blocking=True)
        predictions = best_model(inputs).argmax(dim=1)
        correct_predictions += (predictions == labels).sum().item()
        prediction_count += labels.size(0)

test_accuracy = correct_predictions / prediction_count
print(f"Best checkpoint test accuracy: {test_accuracy:.2%}")

Usaste Ray Tune para programar pruebas simultáneas de GPU, detener configuraciones que no rindieron con ASHA, mantener puntos de control y seleccionar un modelo para la evaluación final. Aumenta NUM_SAMPLES, MAX_EPOCHS, o el tamaño de los subconjuntos de datos cuando quieras una búsqueda más exhaustiva.

Cuaderno de ejemplo

Ajusta un modelo CIFAR-10 con Ray Tune

Obtención del cuaderno

  • Se actualizó por última vez el