Distribua o treinamento PyTorch com o TorchDistributor

6 minutos

O PyTorch, em comum com outras estruturas de aprendizagem profunda como o TensorFlow, foi projetado para ser dimensionado em vários processadores (CPUs ou GPUs) em um único computador. Na maioria dos casos, essa abordagem de expansão usando computadores com processadores mais ou mais rápidos fornece 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 tarefas de processamento em vários nós de trabalho.

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

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

Adapte seu código existente: modifique seu código de treinamento de nó único para ser compatível com o treinamento distribuído. Certifique-se de que sua lógica de treinamento esteja encapsulada em uma única função.
Mover 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 decapagem.
Prepare a função de treinamento: inclua 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 são movidos para o dispositivo apropriado (CPU ou GPU).
Instanciar e executar TorchDistributor: crie uma instância de TorchDistributor com os parâmetros desejados e chame .run(*args) para iniciar o treinamento distribuído.

Adapte seu código existente

Primeiro, você precisa modificar seu código de treinamento de nó único para ser compatível com o treinamento distribuído. Ao modificar seu código, você precisa garantir que sua lógica de treinamento seja encapsulada em uma única função. Esta função é usada pelo TorchDistributor para distribuir o treinamento em 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 seu conjunto de dados que está em um formato compatível com o PyTorch usando torch.utils.data.DataLoadero .

# 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 decapagem, 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 formação

Inclua 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 são 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 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)

Monitorize e avalie o seu trabalho de formação

Você pode usar as ferramentas internas para monitorar o desempenho do cluster, incluindo o uso da CPU ou GPU e a utilização da memória. Quando o treinamento estiver concluído, você poderá avaliar o modelo em um conjunto de dados de validação ou teste usando técnicas de avaliação PyTorch para avaliar o desempenho do seu 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

Gorjeta