Distribución del entrenamiento de PyTorch con TorchDistributor

6 minutos

PyTorch, igual que otros marcos de aprendizaje profundo (como TensorFlow), está diseñado para escalar en varios procesadores (CPU o GPU) en un solo equipo. En la mayoría de los casos, este enfoque para escalar verticalmente mediante equipos con más procesadores o con procesadores más rápidos proporciona un rendimiento de entrenamiento adecuado.

Sin embargo, cuando necesite trabajar con redes neuronales complejas o con grandes volúmenes de datos de entrenamiento, puede beneficiarse de la capacidad inherente de Apache Spark para escalar horizontalmente las tareas de procesamiento en varios nodos de trabajo.

Azure Databricks usa clústeres de Spark que pueden incluir varios nodos de trabajo. Para hacer un uso óptimo de esos clústeres, puede usar TorchDistributor, una biblioteca de código abierto que le permite distribuir trabajos de entrenamiento de PyTorch entre los nodos de un clúster. TorchDistributor está disponible en Databricks Runtime ML 13.0 y versiones posteriores.

Cuando ya ha entrenado un modelo con PyTorch, puede convertir el entrenamiento de proceso único en entrenamiento distribuido con TorchDistributor mediante:

Adaptar el código existente: Modifique el código de entrenamiento de nodo único para que sea compatible con el entrenamiento distribuido. Asegúrese de que la lógica de entrenamiento está encapsulada dentro de una sola función.
Mover importaciones dentro de la función de entrenamiento: Coloque las importaciones necesarias, como import torch, dentro de la función de entrenamiento para evitar errores comunes de selección.
Preparar la función de entrenamiento: Incluya el modelo, el optimizador, la función de pérdida y el bucle de entrenamiento dentro de la función de entrenamiento. Asegúrese de que el modelo y los datos se mueven al dispositivo adecuado (CPU o GPU).
Crear instancias y ejecutar TorchDistributor: Cree una instancia de TorchDistributor con los parámetros deseados y llame a .run(*args) para iniciar el entrenamiento distribuido.

Adaptar el código existente

Primero, modifique el código de entrenamiento de nodo único para que sea compatible con el entrenamiento distribuido. Al modificar el código, debe asegurarse de que la lógica de entrenamiento está encapsulada dentro de una sola función. Esta función la usa TorchDistributor para distribuir el entrenamiento entre varios nodos.

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)

Ahora puede preparar el conjunto de datos que tiene un formato compatible con PyTorch mediante 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 importaciones dentro de la función de entrenamiento

Coloque las importaciones necesarias, como import torch, dentro de la función de entrenamiento para evitar errores comunes de selección. Colocar todas las importaciones dentro de la función de entrenamiento garantiza que todos los módulos necesarios estén disponibles cuando la función se distribuya entre varios nodos.

Preparar la función de entrenamiento

Incluya el modelo, el optimizador, la función de pérdida y el bucle de entrenamiento dentro de la función de entrenamiento. Asegúrese de que el modelo y los datos se mueven al dispositivo adecuado (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()

Crear instancias y ejecutar TorchDistributor

Cree una instancia de TorchDistributor con los parámetros deseados y llame a .run(*args) para iniciar el entrenamiento distribuido. La ejecución de TorchDistributor distribuye las tareas de entrenamiento entre varios nodos.

from pyspark.ml.torch.distributor import TorchDistributor

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

Supervisión y evaluación del trabajo de entrenamiento

Puede usar las herramientas integradas para supervisar el rendimiento del clúster, incluido el uso de CPU o GPU, y el uso de memoria. Una vez completado el entrenamiento, puede evaluar el modelo en un conjunto de datos de validación o prueba mediante técnicas de evaluación de PyTorch para evaluar el rendimiento del 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

Sugerencia