PyTorch - CoreWeave Docs

Utilisez W&B pour le suivi des expériences de machine learning, la gestion des versions des datasets et la collaboration sur vos projets.

Contenu de ce notebook

Ce tutoriel explique comment intégrer W&B à votre code d’entraînement PyTorch pour suivre vos expériences, journaliser des métriques et des gradients, et versionner vos modèles. Utilisez-le pour ajouter le suivi des expériences à un pipeline PyTorch existant.

Schéma de l’intégration de PyTorch et W&B

# importer la bibliothèque
import wandb

# capture a dictionary of hyperparameters with config
config = {
    "learning_rate": 0.001,
    "epochs": 100,
    "batch_size": 128
}

# démarrer une nouvelle expérience
with wandb.init(project="new-sota-model", config=config) as run:

    # configurer le modèle et les données
    model, dataloader = get_model(), get_data()

    # facultatif : suivre les gradients
    run.watch(model)

    for batch in dataloader:
        metrics = model.training_step()
        # journaliser les métriques dans votre boucle d'entraînement pour visualiser les performances du modèle
        run.log(metrics)

    # facultatif : enregistrer le modèle à la fin
    model.to_onnx()
    run.save("model.onnx")

Suivez le tutoriel vidéo en parallèle. Les sections dont le titre commence par Étape contiennent tout ce dont vous avez besoin pour intégrer W&B à un pipeline existant. Le reste du notebook charge les données et définit un modèle.

Installer, importer et se connecter

Avant de définir l’expérience, configurez l’environnement et authentifiez-vous auprès de W&B.

import os
import random

import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from tqdm.auto import tqdm

# Garantir un comportement déterministe
seed = 42
torch.backends.cudnn.deterministic = True
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# Configuration du périphérique
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# retirer le miroir lent de la liste des miroirs MNIST
torchvision.datasets.MNIST.mirrors = [mirror for mirror in torchvision.datasets.MNIST.mirrors
                                      if not mirror.startswith("http://yann.lecun.com")]

Étape 0 : installer W&B

Pour commencer, installez la bibliothèque wandb avec pip.

!pip install wandb onnx -Uq

Étape 1 : importer W&B et se connecter

Pour journaliser des données dans le service W&B, vous devez vous connecter. Si vous utilisez W&B pour la première fois, inscrivez-vous gratuitement en suivant le lien qui s’affiche.

import wandb

wandb.login()

Définir l’expérience et le pipeline

Une fois W&B installé et votre session authentifiée, définissez la configuration de l’expérience ainsi que le pipeline d’entraînement qui l’utilisera.

Suivre les métadonnées et les hyperparamètres avec `wandb.init()`

Commencez par définir votre expérience dans le code. Quels sont les hyperparamètres ? Quelles métadonnées sont associées à ce run ? Un flux de travail courant consiste à stocker ces informations dans un dictionnaire config (ou un objet similaire), puis à y accéder selon les besoins. Cet exemple ne fait varier que quelques hyperparamètres et code le reste en dur. N’importe quel élément de votre modèle peut figurer dans la config. L’exemple inclut également des métadonnées relatives au dataset MNIST et à une architecture convolutive. Si vous travaillez ensuite, par exemple, avec des architectures entièrement connectées sur CIFAR dans le même projet, ces métadonnées vous aideront à distinguer vos runs.

config = dict(
    epochs=5,
    classes=10,
    kernels=[16, 32],
    batch_size=128,
    learning_rate=0.005,
    dataset="MNIST",
    architecture="CNN")

Ensuite, définissez le pipeline global, classique pour l’entraînement d’un modèle :

make : créer un modèle, ainsi que les données et l’optimiseur associés.
train : entraîner le modèle en conséquence.
test : tester le modèle pour évaluer le déroulement de l’entraînement.

Le code suivant implémente ces fonctions.

def model_pipeline(hyperparameters):

    # demander à wandb de démarrer
    with wandb.init(project="pytorch-demo", config=hyperparameters) as run:
        # accéder à tous les hyperparamètres via run.config, pour que la journalisation reflète l’exécution réelle.
        config = run.config

        # créer le modèle, les données et le problème d’optimisation
        model, train_loader, test_loader, criterion, optimizer = make(config)
        print(model)

        # puis les utiliser pour entraîner le modèle
        train(model, train_loader, criterion, optimizer, config)

        # et évaluer ses performances finales
        test(model, test_loader)

    return model

La seule différence avec un pipeline standard est que tout se déroule dans le contexte de wandb.init(). L’appel de cette fonction établit un canal de communication entre votre code et les serveurs W&B. Passer le dictionnaire config à wandb.init() journalise immédiatement toutes ces informations dans W&B : vous savez ainsi toujours quelles valeurs d’hyperparamètres vous avez définies pour votre expérience. Pour garantir que les valeurs choisies et journalisées sont toujours celles qu’utilise votre modèle, W&B recommande d’utiliser la copie run.config de votre objet. Consultez la définition de make ci-dessous pour en voir quelques exemples. Une fois le pipeline défini, les sections suivantes implémentent successivement chacune de ses étapes : configuration des données et du modèle, entraînement et test.

Remarque : W&B exécute son code dans des processus distincts afin qu’un éventuel problème côté W&B ne fasse pas planter votre code. Une fois le problème résolu, vous pouvez journaliser les données avec wandb sync.

def make(config):
    # Préparer les données
    train, test = get_data(train=True), get_data(train=False)
    train_loader = make_loader(train, batch_size=config.batch_size)
    test_loader = make_loader(test, batch_size=config.batch_size)

    # Créer le modèle
    model = ConvNet(config.kernels, config.classes).to(device)

    # Définir la fonction de perte et l’optimiseur
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), lr=config.learning_rate)
    
    return model, train_loader, test_loader, criterion, optimizer

Définir le chargement des données et le modèle

Indiquez ensuite comment les données sont chargées et à quoi ressemble le modèle. Cette étape est importante, mais elle est identique à ce qu’elle serait sans wandb.

def get_data(slice=5, train=True):
    full_dataset = torchvision.datasets.MNIST(root=".",
                                              train=train, 
                                              transform=transforms.ToTensor(),
                                              download=True)
    #  équivalent à un slicing avec [::slice] 
    sub_dataset = torch.utils.data.Subset(
      full_dataset, indices=range(0, len(full_dataset), slice))
    
    return sub_dataset


def make_loader(dataset, batch_size):
    loader = torch.utils.data.DataLoader(dataset=dataset,
                                         batch_size=batch_size, 
                                         shuffle=True,
                                         pin_memory=True, num_workers=2)
    return loader

L’utilisation de wandb ne modifie en rien la définition du modèle ; cet exemple utilise donc une architecture ConvNet standard. N’hésitez pas à expérimenter avec ce code. W&B journalise tous vos résultats sur Forge.

# Réseau de neurones classique et convolutif

class ConvNet(nn.Module):
    def __init__(self, kernels, classes=10):
        super(ConvNet, self).__init__()
        
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, kernels[0], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, kernels[1], kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7 * 7 * kernels[-1], classes)
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

Définir la logique d’entraînement

Poursuivons dans model_pipeline : il est temps de définir le fonctionnement de train. C’est à cette étape que l’intégration W&B suit les gradients, les paramètres et les métriques au fil de l’entraînement. Deux fonctions de wandb entrent ici en jeu : watch et log.

Suivre les gradients avec `run.watch()` et tout le reste avec `run.log()`

run.watch() journalise les gradients et les paramètres de votre modèle toutes les log_freq étapes d’entraînement. Appelez run.watch() avant de lancer l’entraînement. Pour en savoir plus sur les modes de journalisation, l’utilisation de plusieurs modèles et les conseils de performances, consultez Comment puis-je journaliser les gradients et les poids du modèle avec wandb.watch ?. Le reste du code d’entraînement ne change pas : itérez sur les époques et les lots, effectuez les passes avant et arrière, puis appliquez votre optimizer.

def train(model, loader, criterion, optimizer, config):
    # Demander à wandb d’observer le comportement du modèle : gradients, poids, etc.
    run = wandb.init(project="pytorch-demo", config=config)
    run.watch(model, criterion, log="all", log_freq=10)

    # Lancer l’entraînement et le suivre avec wandb
    total_batches = len(loader) * config.epochs
    example_ct = 0  # nombre d’exemples traités
    batch_ct = 0
    for epoch in tqdm(range(config.epochs)):
        for _, (images, labels) in enumerate(loader):

            loss = train_batch(images, labels, model, optimizer, criterion)
            example_ct +=  len(images)
            batch_ct += 1

            # Journaliser les métriques tous les 25 lots
            if ((batch_ct + 1) % 25) == 0:
                train_log(loss, example_ct, epoch)


def train_batch(images, labels, model, optimizer, criterion):
    images, labels = images.to(device), labels.to(device)
    
    # Passe avant ➡
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Passe arrière ⬅
    optimizer.zero_grad()
    loss.backward()

    # Effectuer une étape d’optimisation
    optimizer.step()

    return loss

La seule différence se trouve dans le code de journalisation : là où vous affichiez auparavant les métriques dans le terminal, vous les transmettez désormais à run.log(). run.log() attend un dictionnaire dont les clés sont des chaînes de caractères. Ces chaînes identifient les objets journalisés, qui constituent les valeurs. Vous pouvez également journaliser, de manière facultative, l’étape (step) d’entraînement en cours.

Remarque : utiliser le nombre d’exemples vus par le modèle facilite la comparaison entre différentes tailles de lot, mais vous pouvez aussi utiliser le nombre brut d’étapes ou de lots. Pour les runs d’entraînement plus longs, il peut également être judicieux de journaliser par epoch.

def train_log(loss, example_ct, epoch):
    with wandb.init(project="pytorch-demo") as run:
        # Journaliser la perte et le numéro d’époque
        # C’est ici que les métriques sont journalisées dans W&B
        run.log({"epoch": epoch, "loss": loss}, step=example_ct)
        print(f"Loss after {str(example_ct).zfill(5)} examples: {loss:.3f}")

Définir la logique de test

Une fois l’entraînement du modèle terminé, testez-le : exécutez-le par exemple sur de nouvelles données issues de la production, ou appliquez-le à quelques exemples sélectionnés à la main. L’étape de test est aussi le moment idéal pour enregistrer le modèle entraîné.

Facultatif : appeler `run.save()`

C’est également le bon moment pour enregistrer sur disque l’architecture du modèle et ses paramètres finaux. Pour une compatibilité optimale, exportez (export) le modèle au format Open Neural Network eXchange (ONNX). En transmettant ce nom de fichier à run.save(), vous avez la garantie que les paramètres du modèle sont enregistrés sur les serveurs W&B : fini les incertitudes sur le fichier .h5 ou .pb qui correspond à chaque run d’entraînement. Pour découvrir des fonctionnalités wandb plus avancées de stockage, de gestion des versions et de distribution des modèles, consultez les outils Artifacts.

def test(model, test_loader):
    model.eval()

    with wandb.init(project="pytorch-demo") as run:
        # Exécuter le modèle sur quelques exemples de test
        with torch.no_grad():
            correct, total = 0, 0
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

            print(f"Accuracy of the model on the {total} " +
                f"test images: {correct / total:%}")
            
            run.log({"test_accuracy": correct / total})

        # Enregistrer le modèle au format d’échange ONNX
        torch.onnx.export(model, images, "model.onnx")
        run.save("model.onnx")

Lancer l’entraînement et suivre vos métriques en direct sur wandb.ai

Maintenant que vous avez défini l’ensemble du pipeline et ajouté ces quelques lignes de code W&B, vous pouvez exécuter votre expérience, entièrement suivie. W&B vous indique quelques liens : la documentation, la page du projet (qui regroupe tous les runs d’un projet) et la page du run (où sont stockés les résultats de ce run). Accédez à la page du run et consultez les onglets suivants :

Charts, où les gradients du modèle, les valeurs des paramètres et la perte sont journalisés tout au long de l’entraînement.
System, qui contient les métriques système, notamment l’utilisation des E/S disque ainsi que les métriques CPU et GPU.
Logs, qui contient une copie de tout ce qui a été écrit sur la sortie standard pendant l’entraînement.
Files, où, une fois l’entraînement terminé, vous pouvez cliquer sur model.onnx pour visualiser votre réseau avec le visualiseur de modèles Netron.

Une fois le run terminé, à la sortie du bloc with wandb.init(), W&B affiche également une synthèse des résultats dans la sortie de la cellule.

# Construire, entraîner et analyser le modèle à l’aide du pipeline
model = model_pipeline(config)

Tester des hyperparamètres avec des sweeps

Cet exemple ne portait que sur un seul ensemble d’hyperparamètres. Or, dans la plupart des flux de travail de ML, il est essentiel d’itérer sur plusieurs hyperparamètres. Vous pouvez utiliser W&B Sweeps pour automatiser le test des hyperparamètres et explorer l’espace des modèles et des stratégies d’optimisation possibles. Vous pouvez ainsi aller au-delà du run précédent, limité à une seule configuration. Consultez ce notebook Colab illustrant l’optimisation des hyperparamètres avec W&B Sweeps. L’exécution d’un sweep d’hyperparamètres avec W&B se déroule en trois étapes :

Définir le sweep : créez un dictionnaire ou un fichier YAML qui spécifie les paramètres à explorer, la stratégie de recherche, la métrique d’optimisation, etc.
Initialiser le sweep : sweep_id = wandb.sweep(sweep_config).
Lancer l’agent de sweep : wandb.agent(sweep_id, function=train).

Il n’en faut pas plus pour exécuter un sweep d’hyperparamètres.

Galerie d’exemples

Explorez des exemples de projets suivis et visualisés avec W&B dans la Galerie.

Configuration avancée

Les options suivantes permettent d’étendre le flux de travail de base présenté ci-dessus aux environnements de production, hors ligne ou gérés :

Variables d’environnement : définissez les clés API dans des variables d’environnement pour pouvoir lancer l’entraînement sur un cluster géré.
Mode hors ligne : utilisez le mode dryrun pour entraîner hors ligne et synchroniser les résultats plus tard.
Sur site : installez W&B dans un cloud privé ou sur des serveurs isolés du réseau au sein de votre propre infrastructure.
Sweeps : configurez rapidement une recherche d’hyperparamètres grâce à un outil d’optimisation léger.

​Contenu de ce notebook

​Installer, importer et se connecter

​Étape 0 : installer W&B

​Étape 1 : importer W&B et se connecter

​Définir l’expérience et le pipeline

​Suivre les métadonnées et les hyperparamètres avec wandb.init()

​Définir le chargement des données et le modèle

​Définir la logique d’entraînement

​Suivre les gradients avec run.watch() et tout le reste avec run.log()

​Définir la logique de test

​Facultatif : appeler run.save()

​Lancer l’entraînement et suivre vos métriques en direct sur wandb.ai

​Tester des hyperparamètres avec des sweeps

​Galerie d’exemples

​Configuration avancée