Verteilen von PyTorch-Trainingsaufgaben mit TorchDistributor

• 6 Minuten

PyTorch ist wie andere Deep Learning-Frameworks wie TensorFlow für die Skalierung über mehrere Prozessoren (CPUs oder GPUs) auf einem einzigen Computer ausgelegt. In den meisten Fällen bietet dieser Ansatz zum Hochskalieren mithilfe von Computern mit mehr oder schnelleren Prozessoren eine angemessene Trainingsleistung.

Wenn Sie jedoch mit komplexen neuronalen Netzwerken oder großen Mengen von Trainingsdaten arbeiten müssen, profitieren Sie möglicherweise von der inhärenten Fähigkeit von Apache Spark, Verarbeitungsaufgaben in mehreren Workerknoten aufzuskalieren.

Azure Databricks verwendet Spark-Cluster, die mehrere Workerknoten enthalten können. Um diese Cluster optimal zu nutzen, können Sie TorchDistributor verwenden. Dabei handelt es sich um eine Open-Source-Bibliothek, mit der Sie PyTorch-Trainingsaufträge auf die Knoten in einem Cluster verteilen können. TorchDistributor ist in Databricks Runtime ML 13.0 und höher verfügbar.

Wenn Sie bereits ein Modell mit PyTorch trainiert haben, können Sie Ihr Training mit einem einzelnen Prozess mit TorchDistributor in verteiltes Training umwandeln:

1. Anpassen des vorhandenen Codes: Ändern Sie Ihren Einzelknotentrainingscode so, dass er mit verteiltem Training kompatibel ist. Stellen Sie sicher, dass Ihre Trainingslogik in einer einzelnen Funktion gekapselt ist.
2. Verschieben von Importen innerhalb der Trainingsfunktion: Platzieren Sie erforderliche Importe, wie z. B. import torch, innerhalb der Trainingsfunktion, um häufige Pickle-Fehler zu vermeiden.
3. Vorbereiten der Trainingsfunktion: Schließen Sie Ihr Modell, den Optimierer, die Verlustfunktion und die Trainingsschleife in die Trainingsfunktion ein. Stellen Sie sicher, dass das Modell und die Daten auf das entsprechende Gerät (CPU oder GPU) verschoben werden.
4. Instanziieren und Ausführen von TorchDistributor: Erstellen Sie eine Instanz von TorchDistributor mit den gewünschten Parametern, und rufen Sie .run(*args) auf, um das verteilte Training zu starten.

Anpassen des vorhandenen Codes

Sie müssen zunächst Ihren Einzelknotentrainingscode so ändern, dass er mit verteiltem Training kompatibel ist. Stellen Sie beim Ändern Ihres Codes sicher, dass Ihre Trainingslogik in einer einzelnen Funktion gekapselt ist. Diese Funktion wird von TorchDistributor verwendet, um das Training auf mehrere Knoten zu verteilen.

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)

Jetzt können Sie Ihr Dataset in einem mit PyTorch kompatiblen Format mithilfe von torch.utils.data.DataLoader vorbereiten.

# 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)

Verschieben von Importen innerhalb der Trainingsfunktion

Platzieren Sie erforderliche Importe, wie z. B. import torch, innerhalb der Trainingsfunktion, um häufige Pickle-Fehler zu vermeiden. Durch das Platzieren aller Importe innerhalb der Trainingsfunktion wird sichergestellt, dass alle erforderlichen Module verfügbar sind, wenn die Funktion auf mehrere Knoten verteilt wird.

Vorbereiten der Trainingsfunktion

Schließen Sie Ihr Modell, den Optimierer, die Verlustfunktion und die Trainingsschleife in die Trainingsfunktion ein. Stellen Sie sicher, dass das Modell und die Daten auf das entsprechende Gerät (CPU oder GPU) verschoben werden.

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()

Instanziieren und Ausführen von TorchDistributor

Erstellen Sie eine Instanz von TorchDistributor mit den gewünschten Parametern, und rufen Sie .run(*args) auf, um das verteilte Training zu starten. Bei der Ausführung von TorchDistributor werden die Trainingsaufgaben auf mehrere Knoten verteilt.

from pyspark.ml.torch.distributor import TorchDistributor

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

Überwachen und Bewerten Ihres Trainingsauftrags

Sie können die integrierten Tools verwenden, um die Leistung Ihres Clusters zu überwachen, einschließlich CPU- oder GPU-Auslastung und Arbeitsspeicherauslastung. Wenn das Training abgeschlossen ist, können Sie das Modell anhand eines Validierungs- oder Testdatasets mithilfe von PyTorch-Auswertungstechniken bewerten, um die Leistung Ihres Modells zu beurteilen.

# 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