Distribuera PyTorch-träning med TorchDistributor

  • 6 minuter

PyTorch, i likhet med andra djupinlärningsramverk som TensorFlow, är utformat för att skala över flera processorer (processorer eller GPU:er) på en enda dator. I de flesta fall ger den här metoden för att skala upp med hjälp av datorer med fler eller snabbare processorer tillräcklig träningsprestanda.

Men när du behöver arbeta med komplexa neurala nätverk eller stora mängder träningsdata kan du dra nytta av Apache Sparks inneboende förmåga att skala ut bearbetningsuppgifter över flera arbetsnoder.

Azure Databricks använder Spark-kluster som kan innehålla flera arbetsnoder. Om du vill använda dessa kluster optimalt kan du använda TorchDistributor, ett bibliotek med öppen källkod som gör att du kan distribuera PyTorch-träningsjobb mellan noderna i ett kluster. TorchDistributor är tillgängligt på Databricks Runtime ML 13.0 och senare.

När du redan har tränat en modell med PyTorch kan du konvertera din enda processträning till distribuerad träning med TorchDistributor genom att:

  1. Anpassa din befintliga kod: Ändra träningskoden för en nod så att den är kompatibel med distribuerad träning. Se till att din träningslogik är inkapslad i en enda funktion.
  2. Flytta importer i träningsfunktionen: Placera nödvändiga importer, till exempel import torch, i träningsfunktionen för att undvika vanliga hämtningsfel.
  3. Förbered träningsfunktionen: Inkludera din modell, optimerare, förlustfunktion och träningsloop i träningsfunktionen. Se till att modellen och data flyttas till rätt enhet (CPU eller GPU).
  4. Instansiera och kör TorchDistributor: Skapa en instans av TorchDistributor med önskade parametrar och anropa .run(*args) för att starta den distribuerade träningen.

Anpassa din befintliga kod

Först måste du ändra träningskoden för en nod så att den är kompatibel med distribuerad träning. När du ändrar koden måste du se till att din träningslogik är inkapslad i en enda funktion. Den här funktionen används av TorchDistributor för att distribuera träningen över flera noder.

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)

Nu kan du förbereda din datauppsättning som är i ett format som är kompatibelt med PyTorch med hjälp av 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)

Flytta importer i träningsfunktionen

För att undvika vanliga pickling-fel placerar du nödvändiga importer, till exempel import torch, i träningsfunktionen. Genom att placera alla importer i träningsfunktionen ser du till att alla nödvändiga moduler är tillgängliga när funktionen distribueras över flera noder.

Förbereda träningsfunktionen

Inkludera din modell, optimerare, förlustfunktion och träningsloop i träningsfunktionen. Se till att modellen och data flyttas till rätt enhet (CPU eller GPU).

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

Instansiera och köra TorchDistributor

Skapa en instans av TorchDistributor med önskade parametrar och anropa .run(*args) för att starta den distribuerade träningen. När du kör TorchDistributor distribueras träningsaktiviteterna över flera noder.

from pyspark.ml.torch.distributor import TorchDistributor

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

Övervaka och utvärdera ditt träningsjobb

Du kan använda de inbyggda verktygen för att övervaka klustrets prestanda, inklusive cpu- eller GPU-användning och minnesanvändning. När träningen är klar kan du utvärdera modellen på en validerings- eller testdatauppsättning med hjälp av PyTorch-utvärderingstekniker för att utvärdera modellens prestanda.

# 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

Dricks

Läs mer om distribuerad träning med TorchDistributor.