Distribuir o treinamento do PyTorch com o TorchDistributor

6 minutos

O PyTorch, em comum com outras estruturas de aprendizado profundo, como o TensorFlow, foi projetado para dimensionar em vários processadores (CPUs ou GPUs) em um único computador. Na maioria dos casos, essa abordagem de escala vertical usando computadores com mais processadores ou com processadores mais rápidos fornece um desempenho de treinamento adequado.

No entanto, quando você precisa trabalhar com redes neurais complexas ou grandes volumes de dados de treinamento, você pode se beneficiar da capacidade inerente do Apache Spark de expandir as tarefas de processamento em vários nós de trabalho.

O Azure Databricks usa clusters Spark que podem incluir vários nós de trabalho. Para fazer uso ideal desses clusters, você pode usar o TorchDistributor, uma biblioteca de código aberto que permite distribuir trabalhos de treinamento do PyTorch entre os nós em um cluster. O TorchDistributor está disponível no Databricks Runtime ML 13.0 e superior.

Quando você já treinou um modelo com o PyTorch, você pode converter o seu treinamento de processo único em treinamento distribuído com o TorchDistributor:

Adapte o seu código existente: Modifique o código de treinamento de nó único para ser compatível com o treinamento distribuído. Certifique-se de que a lógica de treinamento esteja encapsulada em uma única função.
Mova as importações dentro da função de treinamento: Coloque as importações necessárias, como import torch, dentro da função de treinamento para evitar erros comuns de seleção.
Prepare a função de treinamento: Inclua o seu modelo, otimizador, função de perda e loop de treinamento dentro da função de treinamento. Certifique-se de que o modelo e os dados sejam movidos para o dispositivo apropriado (CPU ou GPU).
Instancie e execute TorchDistributor: Crie uma instância de TorchDistributor com os parâmetros desejados e chame .run(*args) para iniciar o treinamento distribuído.

Adaptar o seu código existente

Primeiro, você precisa modificar o código de treinamento de nó único para ser compatível com o treinamento distribuído. Ao modificar o seu código, você precisa garantir que a lógica de treinamento esteja encapsulada em uma única função. Esta função é usada por TorchDistributor para distribuir o treinamento entre vários nós.

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)

Agora você pode preparar o conjunto de dados que está em um formato compatível com o 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)

Mover importações dentro da função de treinamento

Para evitar erros comuns de seleção, coloque as importações necessárias, como import torch, dentro da função de treinamento. Colocar todas as importações dentro da função de treinamento garante que todos os módulos necessários estejam disponíveis quando a função for distribuída em vários nós.

Preparar a função de treinamento

Inclua o seu modelo, otimizador, função de perda e loop de treinamento dentro da função de treinamento. Certifique-se de que o modelo e os dados sejam movidos para o dispositivo apropriado (CPU ou 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()

Instanciar e executar o TorchDistributor

Crie uma instância de TorchDistributor com os parâmetros desejados e chame .run(*args) para iniciar o treinamento distribuído. A execução do TorchDistributor distribui as tarefas de treinamento em vários nós.

from pyspark.ml.torch.distributor import TorchDistributor

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

Monitorar e avaliar o seu trabalho de treinamento

Você pode usar as ferramentas internas para monitorar o desempenho do cluster, incluindo uso de CPU ou GPU e utilização de memória. Quando o treinamento for concluído, você poderá avaliar o modelo em um conjunto de dados de validação ou teste usando técnicas de avaliação do PyTorch para avaliar o desempenho do modelo.

# 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

Dica