Distribuer l’entraînement PyTorch avec TorchDistributor

6 minutes

PyTorch, tout comme d’autres infrastructures de Deep Learning comme TensorFlow, est conçu pour être mis à l’échelle sur plusieurs processeurs (unités centrales ou processeurs graphiques) sur un seul ordinateur. Dans la plupart des cas, cette approche du scale-up à l’aide d’ordinateurs dotés de processeurs plus nombreux ou plus rapides offre des performances d’entraînement adéquates.

Toutefois, lorsque vous devez travailler avec des réseaux neuronaux complexes ou de grands volumes de données d’entraînement, vous pouvez tirer parti de la capacité inhérente d’Apache Spark à effectuer un scale-out des tâches de traitement sur plusieurs nœuds worker.

Azure Databricks utilise des clusters Spark qui peuvent inclure plusieurs nœuds worker. Pour optimiser l’utilisation de ces clusters, vous pouvez utiliser TorchDistributor, une bibliothèque open source qui vous permet de distribuer des travaux d’entraînement PyTorch sur les nœuds d’un cluster. TorchDistributor est disponible sur Databricks Runtime ML 13.0 et versions ultérieures.

Lorsque vous avez déjà entraîné un modèle avec PyTorch, vous pouvez convertir votre entraînement de processus unique en entraînement distribué avec TorchDistributor en procédant comme suit :

Adapter votre code existant : Modifiez votre code d’entraînement à nœud unique pour qu’il soit compatible avec l’entraînement distribué. Vérifiez que votre logique d’entraînement est encapsulée dans une seule fonction.
Déplacer les importations dans la fonction d’entraînement : Placez les importations nécessaires, telles que import torch, à l’intérieur de la fonction d’entraînement pour éviter les erreurs de sélection courantes.
Préparer la fonction d’entraînement : Incluez votre modèle, optimiseur, fonction de perte et boucle d’entraînement dans la fonction d’entraînement. Vérifiez que le modèle et les données sont déplacés vers le bon périphérique (processeur ou processeur graphique).
Instancier et exécuter TorchDistributor : Créez une instance de TorchDistributor avec les paramètres souhaités et appelez .run(*args) pour lancer l’entraînement distribué.

Adapter votre code existant

Tout d’abord, vous devez modifier votre code d’entraînement à nœud unique pour qu’il soit compatible avec l’entraînement distribué. Lorsque vous modifiez votre code, vous devez vous assurer que votre logique d’entraînement est encapsulée dans une seule fonction. Cette fonction est utilisée par TorchDistributor pour distribuer l’entraînement sur plusieurs nœuds.

import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.fc(x)

Vous pouvez maintenant préparer votre jeu de données dans un format compatible avec PyTorch à l’aide de torch.utils.data.DataLoader.

# Sample data
inputs = torch.randn(100, 10)
targets = torch.randn(100, 1)

# Create dataset and dataloader
from torch.utils.data import DataLoader, TensorDataset
dataset = TensorDataset(inputs, targets)
dataloader = DataLoader(dataset, batch_size=10)

Déplacer les importations dans la fonction d’entraînement

Pour éviter les erreurs de sélection courantes, placez les importations nécessaires, telles que import torch, à l’intérieur de la fonction d’entraînement. Le placement de toutes les importations dans la fonction d’entraînement garantit que tous les modules requis sont disponibles lorsque la fonction est distribuée sur plusieurs nœuds.

Préparer la fonction d’entraînement

Incluez votre modèle, optimiseur, fonction de perte et boucle d’entraînement dans la fonction d’entraînement. Vérifiez que le modèle et les données sont déplacés vers le bon périphérique (processeur ou processeur graphique).

def train_model(dataloader):
    import torch
    import torch.nn as nn
    from torch.optim import SGD

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SimpleModel().to(device)
    optimizer = SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()
    
    for epoch in range(10):
        for batch in dataloader:
            inputs, targets = batch
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, targets)
            loss.backward()
            optimizer.step()

Instancier et exécuter TorchDistributor

Créez une instance de TorchDistributor avec les paramètres souhaités et appelez .run(*args) pour lancer l’entraînement distribué. L’exécution de TorchDistributor distribue les tâches d’entraînement sur plusieurs nœuds.

from pyspark.ml.torch.distributor import TorchDistributor

# Distribute the training
distributor = TorchDistributor(num_workers=4)
distributor.run(train_model, dataloader)

Surveiller et évaluer votre travail d’entraînement

Vous pouvez utiliser les outils intégrés pour surveiller les performances de votre cluster, notamment l’utilisation du processeur ou du processeur graphique et l’utilisation de la mémoire. Une fois l’entraînement terminé, vous pouvez évaluer le modèle sur un jeu de données de validation ou de test à l’aide de techniques d’évaluation PyTorch pour évaluer les performances de votre modèle.

# Evaluate the model (after distributed training is complete)
model.eval()
with torch.no_grad():
    for inputs, targets in dataloader:
        outputs = model(inputs)
        # Perform evaluation logic

Conseil