PyTorch Ignite - CoreWeave Docs

Cette page explique comment utiliser le handler W&B avec PyTorch Ignite pour journaliser automatiquement les métriques d’entraînement et de validation, les paramètres du modèle et de l’optimiseur, les gradients et les points de contrôle du modèle au cours de vos expériences. Ignite prend en charge un handler W&B qui permet de journaliser les métriques, les paramètres du modèle et de l’optimiseur, ainsi que les gradients pendant l’entraînement et la validation. Vous pouvez également l’utiliser pour journaliser les points de contrôle du modèle dans le cloud W&B. Cette classe encapsule le SDK Python W&B (wandb) ; vous pouvez donc appeler n’importe quelle fonction wandb via ce wrapper. Consultez les exemples qui montrent comment enregistrer les paramètres et les gradients du modèle. Pour aller plus loin, consultez les ressources suivantes :

Voir les visualisations obtenues dans cet exemple de W&B Report.
Essayez d’exécuter le code vous-même dans cet exemple de notebook hébergé.

Configuration de base

L’exemple suivant définit un modèle convolutif simple ainsi que des chargeurs de données pour MNIST. Les exemples de journalisation ci-après s’appuient sur ces éléments.

from argparse import ArgumentParser
import wandb
import torch
from torch import nn
from torch.optim import SGD
from torch.utils.data import DataLoader
import torch.nn.functional as F
from torchvision.transforms import Compose, ToTensor, Normalize
from torchvision.datasets import MNIST

from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Accuracy, Loss

from tqdm import tqdm


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=-1)


def get_data_loaders(train_batch_size, val_batch_size):
    data_transform = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])

    train_loader = DataLoader(MNIST(download=True, root=".", transform=data_transform, train=True),
                              batch_size=train_batch_size, shuffle=True)

    val_loader = DataLoader(MNIST(download=False, root=".", transform=data_transform, train=False),
                            batch_size=val_batch_size, shuffle=False)
    return train_loader, val_loader

L’utilisation de WandBLogger dans Ignite repose sur une approche modulaire. Commencez par créer un objet WandBLogger, puis joignez-le à un trainer ou à un evaluator pour journaliser automatiquement les métriques. Cet exemple montre comment :

Journaliser la perte d’entraînement, en la joignant à l’objet trainer.
Journaliser la perte de validation, en la joignant à l’evaluator.
Journaliser des paramètres facultatifs, comme le taux d’apprentissage.
Surveiller le modèle.

from ignite.contrib.handlers.wandb_logger import *
def run(train_batch_size, val_batch_size, epochs, lr, momentum, log_interval):
    train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size)
    model = Net()
    device = 'cpu'

    if torch.cuda.is_available():
        device = 'cuda'

    optimizer = SGD(model.parameters(), lr=lr, momentum=momentum)
    trainer = create_supervised_trainer(model, optimizer, F.nll_loss, device=device)
    evaluator = create_supervised_evaluator(model,
                                            metrics={'accuracy': Accuracy(),
                                                     'nll': Loss(F.nll_loss)},
                                            device=device)

    desc = "ITERATION - loss: {:.2f}"
    pbar = tqdm(
        initial=0, leave=False, total=len(train_loader),
        desc=desc.format(0)
    )
    #Création de l’objet WandBlogger
    wandb_logger = WandBLogger(
    project="pytorch-ignite-integration",
    name="cnn-mnist",
    config={"max_epochs": epochs,"batch_size":train_batch_size},
    tags=["pytorch-ignite", "mninst"]
    )

    wandb_logger.attach_output_handler(
    trainer,
    event_name=Events.ITERATION_COMPLETED,
    tag="training",
    output_transform=lambda loss: {"loss": loss}
    )

    wandb_logger.attach_output_handler(
    evaluator,
    event_name=Events.EPOCH_COMPLETED,
    tag="training",
    metric_names=["nll", "accuracy"],
    global_step_transform=lambda *_: trainer.state.iteration,
    )

    wandb_logger.attach_opt_params_handler(
    trainer,
    event_name=Events.ITERATION_STARTED,
    optimizer=optimizer,
    param_name='lr'  # facultatif
    )

    wandb_logger.watch(model)

Une fois le logger joint, Ignite transmet automatiquement à votre projet W&B les métriques d’entraînement et de validation, les paramètres de l’optimiseur et les gradients du modèle. Vous pouvez aussi, si vous le souhaitez, utiliser les EVENTS d’Ignite pour journaliser les métriques directement dans le terminal.

    @trainer.on(Events.ITERATION_COMPLETED(every=log_interval))
    def log_training_loss(engine):
        pbar.desc = desc.format(engine.state.output)
        pbar.update(log_interval)

    @trainer.on(Events.EPOCH_COMPLETED)
    def log_training_results(engine):
        pbar.refresh()
        evaluator.run(train_loader)
        metrics = evaluator.state.metrics
        avg_accuracy = metrics['accuracy']
        avg_nll = metrics['nll']
        tqdm.write(
            "Training Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}"
            .format(engine.state.epoch, avg_accuracy, avg_nll)
        )

    @trainer.on(Events.EPOCH_COMPLETED)
    def log_validation_results(engine):
        evaluator.run(val_loader)
        metrics = evaluator.state.metrics
        avg_accuracy = metrics['accuracy']
        avg_nll = metrics['nll']
        tqdm.write(
            "Validation Results - Epoch: {}  Avg accuracy: {:.2f} Avg loss: {:.2f}"
            .format(engine.state.epoch, avg_accuracy, avg_nll))

        pbar.n = pbar.last_print_n = 0

    trainer.run(train_loader, max_epochs=epochs)
    pbar.close()


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument('--batch_size', type=int, default=64,
                        help='input batch size for training (default: 64)')
    parser.add_argument('--val_batch_size', type=int, default=1000,
                        help='input batch size for validation (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10,
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01,
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5,
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--log_interval', type=int, default=10,
                        help='how many batches to wait before logging training status')

    args = parser.parse_args()
    run(args.batch_size, args.val_batch_size, args.epochs, args.lr, args.momentum, args.log_interval)

Ce code génère les visualisations suivantes :

Tableau de bord d'entraînement PyTorch Ignite

Résultats de l'optimisation des hyperparamètres avec PyTorch Ignite

Tableau de bord de comparaison de modèles PyTorch Ignite

Pour plus de détails, référez-vous à la documentation d’Ignite.

​Configuration de base

Configuration de base