Distribuire il training di PyTorch con TorchDistributor

  • 6 minuti

PyTorch, in comune con altri framework di Deep Learning come TensorFlow, è progettato per la scalabilità tra più processori (CPU o GPU) in un singolo computer. Nella maggior parte dei casi, questo approccio alla scalabilità verticale tramite computer con un numero maggiore di processori o processori più veloci offre prestazioni di training adeguate.

Tuttavia, quando è necessario lavorare con reti neurali complesse o con grandi volumi di dati di training, è possibile trarre vantaggio dalla capacità intrinseca di Apache Spark di aumentare le attività di elaborazione su più nodi di lavoro.

Azure Databricks usa cluster Spark che possono includere più nodi di lavoro. Per usare in modo ottimale tali cluster, è possibile usare TorchDistributor, una libreria open source che consente di distribuire i processi di training di PyTorch tra i nodi di un cluster. TorchDistributor è disponibile in Databricks Runtime ML 13.0 e versioni successive.

Quando è già stato eseguito il training di un modello con PyTorch, è possibile convertire il training di un singolo processo in training distribuito con TorchDistributor nei modi seguenti:

  1. Adattare il codice esistente: Modificare il codice di training a nodo singolo in modo che sia compatibile con il training distribuito. Assicurarsi che la logica di training sia incapsulata all'interno di una singola funzione.
  2. Spostare le importazioni all'interno della funzione di training: Inserire le importazioni necessarie, ad esempio import torch, all'interno della funzione di training per evitare errori di pickling comuni.
  3. Preparare la funzione di training: Includere il modello, l'ottimizzatore, la funzione di perdita e il ciclo di training all'interno della funzione di training. Assicurarsi che il modello e i dati vengano spostati nel dispositivo appropriato (CPU o GPU).
  4. Creare un'istanza ed eseguire TorchDistributor: Creare un'istanza di TorchDistributor con i parametri desiderati e chiamare .run(*args) per avviare il training distribuito.

Adattare il codice esistente

È prima di tutto necessario modificare il codice di training a nodo singolo in modo che sia compatibile con il training distribuito. Quando si modifica il codice, è necessario assicurarsi che la logica di training sia incapsulata all'interno di una singola funzione. Questa funzione viene usata da TorchDistributor per distribuire il training tra più nodi.

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)

È ora possibile preparare il set di dati in un formato compatibile con PyTorch usando 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)

Spostare le importazioni all'interno della funzione di training

Per evitare errori di pickling comuni, inserire le importazioni necessarie, ad esempio import torch, all'interno della funzione di training. L'inserimento di tutte le importazioni all'interno della funzione di training garantisce che tutti i moduli necessari siano disponibili quando la funzione viene distribuita tra più nodi.

Preparare la funzione di training

Includere il modello, l'ottimizzatore, la funzione di perdita e il ciclo di training all'interno della funzione di training. Assicurarsi che il modello e i dati vengano spostati nel dispositivo appropriato (CPU o 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()

Creare un'istanza ed eseguire TorchDistributor

Creare un'istanza di TorchDistributor con i parametri desiderati e chiamare .run(*args) per avviare il training distribuito. L'esecuzione di TorchDistributor distribuisce le attività di training tra più nodi.

from pyspark.ml.torch.distributor import TorchDistributor

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

Monitorare e valutare il processo di training

È possibile usare gli strumenti predefiniti per monitorare le prestazioni del cluster, tra cui l'utilizzo della CPU o della GPU e l'utilizzo della memoria. Al termine del training, è possibile valutare il modello in un set di dati di convalida o di test usando tecniche di valutazione di PyTorch per valutare le prestazioni del modello.

# 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

Suggerimento

Altre informazioni sul training distribuito con TorchDistributor.