PyTorch Lightning

PyTorch Lightning fournit un wrapper léger pour organiser votre code PyTorch et ajouter des fonctionnalités avancées telles que l’entraînement distribué et la précision en 16 bits. W&B fournit un wrapper léger pour consigner vos expériences de ML. Vous n’avez pas besoin de combiner les deux vous-même : la bibliothèque PyTorch Lightning inclut directement W&B via le WandbLogger. Cette page explique comment utiliser WandbLogger pour suivre les métriques, consigner les hyperparamètres, enregistrer des points de contrôle du modèle comme artifacts, consigner des médias et exécuter un entraînement multi-GPU avec PyTorch Lightning et W&B.

Intégrer Lightning

Les sections suivantes montrent comment vous authentifier auprès de W&B, installer la bibliothèque wandb et joindre un WandbLogger à votre instance Lightning Trainer ou Fabric.

Logger PyTorch
Logger Fabric

from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch import Trainer

wandb_logger = WandbLogger(log_model="all")
trainer = Trainer(logger=wandb_logger)

Utilisation de wandb.log() : Le WandbLogger envoie les journaux vers W&B en utilisant le global_step du Trainer. Si vous effectuez des appels supplémentaires à wandb.log() directement dans votre code, n’utilisez pas l’argument step dans wandb.log().À la place, journalisez le global_step du Trainer comme vos autres métriques :

wandb.log({"accuracy":0.99, "trainer/global_step": step})

import lightning as L
from wandb.integration.lightning.fabric import WandbLogger

wandb_logger = WandbLogger(log_model="all")
fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({"important_metric": important_metric})

Une clé API permet d’authentifier votre machine auprès de W&B. Vous pouvez générer une clé API depuis votre profil.

Pour une méthode plus directe, accédez aux Paramètres utilisateur et créez une clé API. Copiez immédiatement la clé API et conservez-la dans un endroit sûr, par exemple dans un gestionnaire de mots de passe.

Pour générer une clé API depuis votre profil :

Cliquez sur l’icône de votre profil en haut à droite.
Sélectionnez Paramètres utilisateur, puis faites défiler jusqu’à la section Clés API.

Installez la bibliothèque `wandb` et connectez-vous

Pour installer la bibliothèque wandb localement et vous connecter :

Ligne de commande
Python
Python notebook

Définissez la variable d’environnement WANDB_API_KEY sur votre clé API. Remplacez les valeurs entre <> par les vôtres :
export WANDB_API_KEY=<your_api_key>
Installez la bibliothèque wandb et connectez-vous.
pip install wandb wandb login

pip install wandb

import wandb
wandb.login()

!pip install wandb

import wandb
wandb.login()

Utiliser le `WandbLogger` de PyTorch Lightning

PyTorch Lightning propose plusieurs classes WandbLogger pour consigner des métriques, ainsi que les poids du modèle et des médias. Choisissez la classe qui correspond à votre configuration d’entraînement :

Pour l’utiliser avec Lightning, instanciez WandbLogger, puis transmettez-le à Trainer ou à Fabric de Lightning.

Logger PyTorch
Logger Fabric

trainer = Trainer(logger=wandb_logger)

fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({
    "important_metric": important_metric
})

Arguments courants du logger

Le tableau suivant répertorie les paramètres courants de WandbLogger. Consultez la documentation de PyTorch Lightning pour en savoir plus sur tous les arguments du logger.

Paramètre	Description
`project`	Définit dans quel projet W&B journaliser
`name`	Donne un nom à votre run W&B
`log_model`	Journalise tous les modèles si `log_model="all"`, ou à la fin de l’entraînement si `log_model=True`
`save_dir`	Chemin où les données sont enregistrées

Journalisez vos hyperparamètres

La journalisation des hyperparamètres avec W&B vous permet de comparer les runs et de reproduire les résultats. Utilisez la méthode correspondante à votre logger :

Logger PyTorch
Logger Fabric

class LitModule(LightningModule):
    def __init__(self, *args, **kwarg):
        self.save_hyperparameters()

wandb_logger.log_hyperparams(
    {
        "hyperparameter_1": hyperparameter_1,
        "hyperparameter_2": hyperparameter_2,
    }
)

Journaliser des paramètres de configuration supplémentaires

Pour capturer des valeurs de configuration supplémentaires en plus de vos hyperparamètres, mettez directement à jour la configuration du run :

# ajouter un paramètre
wandb_logger.experiment.config["key"] = value

# ajouter plusieurs paramètres
wandb_logger.experiment.config.update({key1: val1, key2: val2})

# utiliser directement le module wandb
wandb.config["key"] = value
wandb.config.update()

Journaliser les gradients, l’histogramme des paramètres et la topologie du modèle

Passez l’objet de votre modèle à wandb_logger.watch() pour surveiller les gradients et les paramètres de votre modèle pendant l’entraînement. Voir la documentation de WandbLogger pour PyTorch Lightning.

Journaliser des métriques

Logger PyTorch
Fabric Logger

Pour journaliser vos métriques dans W&B lorsque vous utilisez WandbLogger, appelez self.log('my_metric_name', metric_value) dans votre LightningModule, par exemple dans les méthodes training_step ou validation_step.L’extrait de code suivant montre comment définir votre LightningModule pour journaliser vos métriques ainsi que les hyperparamètres de votre LightningModule. Cet exemple utilise la bibliothèque torchmetrics pour calculer vos métriques.

import torch
from torch.nn import Linear, CrossEntropyLoss, functional as F
from torch.optim import Adam
from torchmetrics.functional import accuracy
from lightning.pytorch import LightningModule


class My_LitModule(LightningModule):
    def __init__(self, n_classes=10, n_layer_1=128, n_layer_2=256, lr=1e-3):
        """méthode utilisée pour définir les paramètres du modèle"""
        super().__init__()

        # les images MNIST ont la forme (1, 28, 28) (canaux, largeur, hauteur)
        self.layer_1 = Linear(28 * 28, n_layer_1)
        self.layer_2 = Linear(n_layer_1, n_layer_2)
        self.layer_3 = Linear(n_layer_2, n_classes)

        self.loss = CrossEntropyLoss()
        self.lr = lr

        # enregistrer les hyperparamètres dans self.hparams (journalisés automatiquement par W&B)
        self.save_hyperparameters()

    def forward(self, x):
        """méthode utilisée pour l’inférence, de l’entrée à la sortie"""

        # (b, 1, 28, 28) -> (b, 1*28*28)
        batch_size, channels, width, height = x.size()
        x = x.view(batch_size, -1)

        # appliquer 3 x (linear + relu)
        x = F.relu(self.layer_1(x))
        x = F.relu(self.layer_2(x))
        x = self.layer_3(x)
        return x

    def training_step(self, batch, batch_idx):
        """doit renvoyer une perte à partir d’un seul lot"""
        _, loss, acc = self._get_preds_loss_accuracy(batch)

        # Journaliser la perte et la métrique
        self.log("train_loss", loss)
        self.log("train_accuracy", acc)
        return loss

    def validation_step(self, batch, batch_idx):
        """utilisée pour journaliser des métriques"""
        preds, loss, acc = self._get_preds_loss_accuracy(batch)

        # Journaliser la perte et la métrique
        self.log("val_loss", loss)
        self.log("val_accuracy", acc)
        return preds

    def configure_optimizers(self):
        """définit l’optimiseur du modèle"""
        return Adam(self.parameters(), lr=self.lr)

    def _get_preds_loss_accuracy(self, batch):
        """fonction utilitaire, car les étapes d’entraînement, de validation et de test sont similaires"""
        x, y = batch
        logits = self(x)
        preds = torch.argmax(logits, dim=1)
        loss = self.loss(logits, y)
        acc = accuracy(preds, y)
        return preds, loss, acc

import lightning as L
import torch
import torchvision as tv
from wandb.integration.lightning.fabric import WandbLogger
import wandb

fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()

model = tv.models.resnet18()
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
model, optimizer = fabric.setup(model, optimizer)

train_dataloader = fabric.setup_dataloaders(
    torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
)

model.train()
for epoch in range(num_epochs):
    for batch in train_dataloader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()
        fabric.log_dict({"loss": loss})

Journaliser les valeurs min/max d’une métrique

Avec la fonction define_metric de W&B, vous pouvez définir si votre métrique de synthèse W&B affiche la valeur minimale, maximale, moyenne ou la meilleure valeur pour cette métrique. Si define_metric n’est pas utilisé, la dernière valeur enregistrée apparaît dans vos métriques de synthèse. Pour plus d’informations, voir le guide de personnalisation des axes de journalisation. Pour suivre la précision de validation maximale dans la métrique de synthèse W&B, appelez wandb.define_metric() une seule fois, au début de l’entraînement :

Logger PyTorch
Logger Fabric

class My_LitModule(LightningModule):
    ...

    def validation_step(self, batch, batch_idx):
        if trainer.global_step == 0:
            wandb.define_metric("val_accuracy", summary="max")

        preds, loss, acc = self._get_preds_loss_accuracy(batch)

        # Journaliser la perte et la métrique
        self.log("val_loss", loss)
        self.log("val_accuracy", acc)
        return preds

wandb.define_metric("val_accuracy", summary="max")
fabric = L.Fabric(loggers=[wandb_logger])
fabric.launch()
fabric.log_dict({"val_accuracy": val_accuracy})

Créer un point de contrôle du modèle

L’enregistrement de points de contrôle du modèle en tant qu’artifacts W&B vous fournit des fichiers de modèle versionnés que vous pourrez récupérer ultérieurement par run, alias ou version. Pour enregistrer des points de contrôle du modèle en tant qu’Artifacts W&B Artifacts, utilisez le callback Lightning ModelCheckpoint et définissez l’argument log_model dans WandbLogger.

Logger PyTorch
Logger Fabric

trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback])

fabric = L.Fabric(loggers=[wandb_logger], callbacks=[checkpoint_callback])

Les alias latest et best sont définis automatiquement pour vous permettre de récupérer facilement un point de contrôle du modèle depuis un artifact W&B :

# la référence peut être récupérée dans le panneau des artifacts
# <version> peut être une version (par exemple, "v2") ou un alias ("latest" ou "best")
checkpoint_reference = "<user>/<project>/<model-run_id>:<version>"

Via Logger
Via wandb

# télécharger le point de contrôle du modèle en local (s'il n'est pas déjà en cache)
wandb_logger.download_artifact(checkpoint_reference, artifact_type="model")

# télécharger le point de contrôle du modèle en local (s'il n'est pas déjà en cache)
run = wandb.init(project="MNIST")
artifact = run.use_artifact(checkpoint_reference, type="model")
artifact_dir = artifact.download()

Logger PyTorch
Logger Fabric

# charger le point de contrôle du modèle
model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt")

# récupérer le point de contrôle du modèle brut
full_checkpoint = fabric.load(Path(artifact_dir) / "model.ckpt")

model.load_state_dict(full_checkpoint["model"])
optimizer.load_state_dict(full_checkpoint["optimizer"])

Les points de contrôle du modèle que vous journalisez sont visibles dans l’UI W&B Artifacts et incluent la traçabilité complète du modèle (voir un exemple de point de contrôle du modèle dans l’UI). Pour mettre vos meilleurs points de contrôle du modèle en favoris et les centraliser pour votre équipe, liez-les au registre de modèles W&B. Dans le registre, vous pouvez organiser vos meilleurs modèles par tâche, gérer le cycle de vie des modèles, faciliter leur suivi et leur audit tout au long du cycle de vie du ML, et automatiser les actions en aval avec des webhooks ou des jobs.

Journaliser des images, du texte et plus encore

Le WandbLogger fournit les méthodes log_image, log_text et log_table pour la journalisation de médias. Vous pouvez aussi appeler directement wandb.log() ou trainer.logger.experiment.log() pour journaliser d’autres types de médias, comme Audio, des molécules, des nuages de points et des objets 3D.

Journaliser des images
Journaliser du texte
Journaliser des tableaux

# utiliser des tenseurs, des tableaux NumPy ou des images PIL
wandb_logger.log_image(key="samples", images=[img1, img2])

# ajouter des légendes
wandb_logger.log_image(key="samples", images=[img1, img2], caption=["tree", "person"])

# utiliser un chemin de fichier
wandb_logger.log_image(key="samples", images=["img_1.jpg", "img_2.jpg"])

# utiliser .log dans le trainer
trainer.logger.experiment.log(
    {"samples": [wandb.Image(img, caption=caption) for (img, caption) in my_images]},
    step=current_trainer_global_step,
)

# les données doivent être une liste de listes
columns = ["input", "label", "prediction"]
my_data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]]

# utiliser les colonnes et les données
wandb_logger.log_text(key="my_samples", columns=columns, data=my_data)

# utiliser un DataFrame pandas
wandb_logger.log_text(key="my_samples", dataframe=my_dataframe)

# journaliser un tableau W&B contenant une légende textuelle, une image et de l'audio
columns = ["caption", "image", "sound"]

# les données doivent être une liste de listes
my_data = [
    ["cheese", wandb.Image(img_1), wandb.Audio(snd_1)],
    ["wine", wandb.Image(img_2), wandb.Audio(snd_2)],
]

# journaliser le tableau
wandb_logger.log_table(key="my_samples", columns=columns, data=my_data)

Utilisez le système de callbacks de Lightning pour contrôler quand vous journalisez vers W&B via le WandbLogger. L’exemple suivant journalise un échantillon d’images de validation et de prédictions :

import torch
import wandb
import lightning.pytorch as pl
from lightning.pytorch.loggers import WandbLogger

# or
# from wandb.integration.lightning.fabric import WandbLogger


class LogPredictionSamplesCallback(Callback):
    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
    ):
        """Called when the validation batch ends."""

        # `outputs` comes from `LightningModule.validation_step`
        # which corresponds to our model predictions in this case

        # Journaliser 20 exemples de prédictions d'images du premier batch
        if batch_idx == 0:
            n = 20
            x, y = batch
            images = [img for img in x[:n]]
            captions = [
                f"Ground Truth: {y_i} - Prediction: {y_pred}"
                for y_i, y_pred in zip(y[:n], outputs[:n])
            ]

            # Option 1: log images with `WandbLogger.log_image`
            wandb_logger.log_image(key="sample_images", images=images, caption=captions)

            # Option 2: log images and predictions as a W&B Table
            columns = ["image", "ground truth", "prediction"]
            data = [
                [wandb.Image(x_i), y_i, y_pred] or x_i,
                y_i,
                y_pred in list(zip(x[:n], y[:n], outputs[:n])),
            ]
            wandb_logger.log_table(key="sample_table", columns=columns, data=data)


trainer = pl.Trainer(callbacks=[LogPredictionSamplesCallback()])

Utiliser plusieurs GPU avec Lightning et W&B

Lorsque vous exécutez un entraînement distribué, la manière dont vous référencez wandb.run d’un rang à l’autre peut déterminer si l’entraînement se poursuit ou s’il entre en interblocage. Cette section explique les exigences et présente un modèle recommandé. PyTorch Lightning prend en charge le multi-GPU via son interface DDP. Cependant, la conception de PyTorch Lightning exige que vous fassiez attention à la façon dont vous instanciez vos GPU. Lightning exige que chaque GPU (ou rang) dans votre boucle d’entraînement soit instancié exactement de la même manière, avec les mêmes conditions initiales. Cependant, seul le processus de rang 0 a accès à l’objet wandb.run. Pour les processus de rang non nul, wandb.run = None. Cela peut faire échouer vos processus non nuls. Une telle situation peut vous placer dans un interblocage, car le processus de rang 0 attend que les processus de rang non nul le rejoignent, alors qu’ils ont déjà planté. Pour cette raison, faites attention à la façon dont vous configurez votre code d’entraînement. La méthode recommandée consiste à rendre votre code indépendant de l’objet wandb.run.

class MNISTClassifier(pl.LightningModule):
    def __init__(self):
        super(MNISTClassifier, self).__init__()

        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

        self.loss = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = self.loss(y_hat, y)

        self.log("train/loss", loss)
        return {"train_loss": loss}

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.forward(x)
        loss = self.loss(y_hat, y)

        self.log("val/loss", loss)
        return {"val_loss": loss}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)


def main():
    # Définir toutes les graines aléatoires sur la même valeur.
    # C'est important dans un contexte d'entraînement distribué.
    # Chaque rang obtiendra son propre ensemble de poids initiaux.
    # S'ils ne correspondent pas, les gradients ne correspondront pas non plus,
    # ce qui risque de mener à un entraînement qui ne converge pas.
    pl.seed_everything(1)

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)

    model = MNISTClassifier()
    wandb_logger = WandbLogger(project="<project-name>")
    callbacks = [
        ModelCheckpoint(
            dirpath="checkpoints",
            every_n_train_steps=100,
        ),
    ]
    trainer = pl.Trainer(
        max_epochs=3, gpus=2, logger=wandb_logger, strategy="ddp", callbacks=callbacks
    )
    trainer.fit(model, train_loader, val_loader)

Exemples

Pour un guide pas à pas complet, vous pouvez suivre ce tutoriel dans une vidéo avec un notebook Colab.

Foire aux questions

Comment W&B s’intègre-t-il à Lightning ?

L’intégration principale s’appuie sur l’API Lightning loggers, qui vous permet d’écrire une grande partie de votre code de journalisation de manière indépendante du framework. Les instances de Logger sont transmises au Lightning Trainer et sont déclenchées par le riche système de hooks et de callbacks de cette API. Cela permet de bien séparer votre code de recherche du code d’ingénierie et de journalisation.

Que l’intégration enregistre-t-elle sans code supplémentaire ?

W&B enregistre vos points de contrôle du modèle, que vous pouvez consulter ou télécharger pour les utiliser dans de futures exécutions. W&B capture également les métriques système, comme l’utilisation du GPU et les E/S réseau. Il capture des informations sur l’environnement, telles que les informations sur le matériel et le système d’exploitation. Il capture l’état du code, y compris le commit Git et le patch de diff, le contenu du notebook et l’historique de session. Il capture également tout ce qui est affiché sur la sortie standard.

Que faire si j’ai besoin d’utiliser `wandb.run` dans ma configuration d’entraînement ?

Vous devez vous-même élargir la portée de la variable à laquelle vous souhaitez accéder. En d’autres termes, veillez à ce que les conditions initiales soient identiques dans tous les processus.

if os.environ.get("LOCAL_RANK", None) is None:
    os.environ["WANDB_DIR"] = wandb.run.dir

Si c’est le cas, vous pouvez utiliser os.environ["WANDB_DIR"] pour configurer le répertoire des points de contrôle du modèle. Ainsi, tout processus de rang non nul peut accéder à wandb.run.dir.

Guides

Intégrations

Référence

Intégrer Lightning

Installez la bibliothèque `wandb` et connectez-vous

Utiliser le `WandbLogger` de PyTorch Lightning

Arguments courants du logger

Journalisez vos hyperparamètres

Journaliser des paramètres de configuration supplémentaires

Journaliser les gradients, l’histogramme des paramètres et la topologie du modèle

Journaliser des métriques

Journaliser les valeurs min/max d’une métrique

Créer un point de contrôle du modèle

Journaliser des images, du texte et plus encore

Utiliser plusieurs GPU avec Lightning et W&B

Exemples

Foire aux questions

Comment W&B s’intègre-t-il à Lightning ?

Que l’intégration enregistre-t-elle sans code supplémentaire ?

Que faire si j’ai besoin d’utiliser `wandb.run` dans ma configuration d’entraînement ?

​Intégrer Lightning

​Inscrivez-vous et créez une clé API

​Installez la bibliothèque wandb et connectez-vous

​Utiliser le WandbLogger de PyTorch Lightning

​Arguments courants du logger

​Journalisez vos hyperparamètres

​Journaliser des paramètres de configuration supplémentaires

​Journaliser les gradients, l’histogramme des paramètres et la topologie du modèle

​Journaliser des métriques

​Journaliser les valeurs min/max d’une métrique

​Créer un point de contrôle du modèle

​Journaliser des images, du texte et plus encore

​Utiliser plusieurs GPU avec Lightning et W&B

​Exemples

​Foire aux questions

​Comment W&B s’intègre-t-il à Lightning ?

​Que l’intégration enregistre-t-elle sans code supplémentaire ?

​Que faire si j’ai besoin d’utiliser wandb.run dans ma configuration d’entraînement ?

Intégrer Lightning

Inscrivez-vous et créez une clé API

Installez la bibliothèque `wandb` et connectez-vous

Utiliser le `WandbLogger` de PyTorch Lightning

Arguments courants du logger

Journalisez vos hyperparamètres

Journaliser des paramètres de configuration supplémentaires

Journaliser les gradients, l’histogramme des paramètres et la topologie du modèle

Journaliser des métriques

Journaliser les valeurs min/max d’une métrique

Créer un point de contrôle du modèle

Journaliser des images, du texte et plus encore

Utiliser plusieurs GPU avec Lightning et W&B

Exemples

Foire aux questions

Comment W&B s’intègre-t-il à Lightning ?

Que l’intégration enregistre-t-elle sans code supplémentaire ?

Que faire si j’ai besoin d’utiliser `wandb.run` dans ma configuration d’entraînement ?