PyTorch-training distribueren met TorchDistributor

  • 6 minuten

PyTorch, in combinatie met andere deep learning-frameworks zoals TensorFlow, is ontworpen om te schalen op meerdere processors (CPU's of GPU's) op één computer. In de meeste gevallen biedt deze aanpak om omhoog te schalen door computers met meer of snellere processors te gebruiken, voldoende trainingsprestaties.

Wanneer u echter moet werken met complexe neurale netwerken of grote hoeveelheden trainingsgegevens, kunt u profiteren van de inherente mogelijkheid van Apache Spark om verwerkingstaken uit te schalen op meerdere werkknooppunten.

Azure Databricks maakt gebruik van Spark-clusters die meerdere werkknooppunten kunnen bevatten. Als u optimaal gebruik wilt maken van deze clusters, kunt u TorchDistributor gebruiken, een opensource-bibliotheek waarmee u PyTorch-trainingstaken over de knooppunten in een cluster kunt distribueren. TorchDistributor is beschikbaar op Databricks Runtime ML 13.0 en hoger.

Wanneer u al een model hebt getraind met PyTorch, kunt u uw training voor één proces converteren naar gedistribueerde training met TorchDistributor door:

  1. Pas uw bestaande code aan: pas uw trainingscode met één knooppunt aan zodat deze compatibel is met gedistribueerde training. Zorg ervoor dat uw trainingslogica is ingekapseld binnen één functie.
  2. Importbewerkingen verplaatsen binnen de trainingsfunctie: plaats de benodigde importbewerkingen, zoals import torch, in de trainingsfunctie om veelvoorkomende picklingfouten te voorkomen.
  3. De trainingsfunctie voorbereiden: Neem uw model, optimalisatie, verliesfunctie en trainingslus op in de trainingsfunctie. Zorg ervoor dat het model en de gegevens naar het juiste apparaat (CPU of GPU) worden verplaatst.
  4. Instantieer en voer TorchDistributor uit: Maak een exemplaar van TorchDistributor met de gewenste parameters en roep .run(*args) aan om de gedistribueerde training te starten.

Uw bestaande code aanpassen

Eerst moet u uw trainingscode met één knooppunt wijzigen zodat deze compatibel is met gedistribueerde training. Wanneer u de code wijzigt, moet u ervoor zorgen dat uw trainingslogica is ingekapseld binnen één functie. Deze functie wordt gebruikt door TorchDistributor om de training over meerdere knooppunten te verdelen.

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)

U kunt nu uw gegevensset voorbereiden in een indeling die compatibel is met PyTorch met behulp van 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)

Importbewerkingen verplaatsen binnen de trainingsfunctie

Als u veelvoorkomende kiezerfouten wilt voorkomen, plaatst u de benodigde importbewerkingen, zoals import torch, in de trainingsfunctie. Als u alle importbewerkingen binnen de trainingsfunctie plaatst, zorgt u ervoor dat alle vereiste modules beschikbaar zijn wanneer de functie over meerdere knooppunten wordt verdeeld.

De trainingsfunctie voorbereiden

Neem uw model, optimizer, verliesfunctie en trainingslus op in de trainingsfunctie. Zorg ervoor dat het model en de gegevens naar het juiste apparaat (CPU of GPU) worden verplaatst.

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

TorchDistributor instantiëren en uitvoeren

Maak een exemplaar van TorchDistributor met de gewenste parameters en roep .run(*args) aan om de gedistribueerde training te starten. Door TorchDistributor uit te voeren, worden de trainingstaken verdeeld over meerdere knooppunten.

from pyspark.ml.torch.distributor import TorchDistributor

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

Uw trainingstaak bewaken en evalueren

U kunt de ingebouwde hulpprogramma's gebruiken om de prestaties van uw cluster te bewaken, inclusief CPU- of GPU-gebruik en geheugengebruik. Wanneer de training is voltooid, kunt u het model evalueren op een validatie- of testgegevensset met behulp van PyTorch-evaluatietechnieken om de prestaties van uw model te beoordelen.

# 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

Tip

Meer informatie over gedistribueerde training met TorchDistributor.