En este tutorial, aprenderemos a usar nn.parallel.DistributedDataParallel para entrenar nuestros modelos en múltiples GPU. Tomaremos un ejemplo mínimo de entrenamiento de un clasificador de imágenes y veremos cómo podemos acelerar el entrenamiento.
Comencemos con algunas importaciones.
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import time
Usaremos el CIFAR10 en todos nuestros experimentos con un tamaño de lote de 256.
def create_data_loader_cifar10():
transform = transforms.Compose(
[
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 256
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=10, pin_memory=True)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=10)
return trainloader, testloader
Primero entrenaremos el modelo en una sola GPU Nvidia A100 durante 1 época. Cosas estándar de pytorch aquí, nada nuevo. El tutorial se basa en el tutorial oficial de los documentos de Pytorch.
def train(net, trainloader):
print("Start training...")
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
epochs = 1
num_of_batches = len(trainloader)
for epoch in range(epochs):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
images, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'[Epoch /] loss: ')
print('Finished Training')
El test la función se define de manera similar. El script principal simplemente pondrá todo junto:
if __name__ == '__main__':
start = time.time()
PATH = './cifar_net.pth'
trainloader, testloader = create_data_loader_cifar10()
net = torchvision.models.resnet50(False).cuda()
start_train = time.time()
train(net, trainloader)
end_train = time.time()
torch.save(net.state_dict(), PATH)
test(net, PATH, testloader)
end = time.time()
seconds = (end - start)
seconds_train = (end_train - start_train)
print(f"Total elapsed time: seconds,
Train 1 epoch seconds")
Usamos un resnet50 para medir el rendimiento de una red de tamaño decente.
Ahora vamos a entrenar el modelo:
$ python -m train_1gpu
Accuracy of the network on the 10000 test images: 27 %
Total elapsed time: 69.03 seconds, Train 1 epoch 13.08 seconds
Bien, es hora de comenzar el trabajo de optimización.
El código está disponible en GitHub. Si planea solidificar su conocimiento de Pytorch, hay dos libros increíbles que recomendamos encarecidamente: Aprendizaje profundo con PyTorch de Publicaciones Manning y Aprendizaje automático con PyTorch y Scikit-Learn de Sebastián Raschka. Siempre puedes usar el código de descuento del 35% blaissummer21 para todos los productos de Manning.
torch.nn.DataParallel: sin dolor, sin ganancia
DataParallel es de un solo proceso, multiproceso y solo funciona en una sola máquina. Para cada GPU, usamos el mismo modelo para hacer el pase hacia adelante. Repartimos los datos por las GPU y realizamos pases de avance en cada una de ellas. Esencialmente, lo que sucede es que el tamaño del lote se divide entre la cantidad de trabajadores.
En este caso de uso, esta funcionalidad no proporcionó ninguna ganancia. Eso es porque el sistema que estoy usando tiene un cuello de botella en la CPU y en el disco duro. Otras máquinas que tienen un disco y una CPU muy rápidos pero que luchan con la velocidad de la GPU (cuello de botella de la GPU) pueden beneficiarse de esta funcionalidad.
En la práctica, el único cambio que necesita hacer en el código es el siguiente:
net = torchvision.models.resnet50(False)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
net = nn.DataParallel(net)
Cuando usas nn.DataParallelel tamaño del lote debe ser divisible por el número de GPU.
nn.DataParalleldivide el lote y lo procesa de forma independiente en todas las GPU disponibles. En cada paso hacia adelante, el módulo se replica en cada GPU, lo que representa una sobrecarga significativa. Cada réplica maneja una parte del lote (batch_size/gpus). Durante el paso hacia atrás, los gradientes de cada réplica se suman en el módulo original.
Más información en nuestro artículo anterior sobre paralelismo de datos vs modelo.
Una buena práctica al usar varias GPU es definir de antemano las GPU que usará su secuencia de comandos:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
Esto debe hacerse antes cualquier otra importación relacionada con CUDA.
Incluso del Pytorch documentación es obvio que esta es una estrategia muy pobre:
Se recomienda utilizar
nn.DistributedDataParallelen lugar de esta clase, para realizar un entrenamiento multi-GPU, incluso si solo hay un solo nodo.
La razón es que DistributedDataParallel usa un proceso por trabajador (GPU) mientras que DataParallel encapsula toda la comunicación de datos en un solo proceso.
Según los documentos, los datos pueden estar en cualquier dispositivo antes de pasar al modelo.
En mi experimento, DataParallel fue Más lento que entrenar en una sola GPU. Incluso con 4 GPU. Después de aumentar la cantidad de trabajadores, reduje el tiempo, pero aún peor que una sola GPU. Mido e informo el tiempo requerido para entrenar el modelo para una época, es decir, imágenes de 50K 32×32.
Nota final: para comparar el rendimiento con una sola GPU, multipliqué el tamaño del lote por la cantidad de trabajadores, es decir, 4 por 4 GPU. De lo contrario, es más de 2 veces más lento.
Esto nos lleva al tema principal de Distributed Data-Parallel.
El código está disponible en GitHub. Siempre puede apoyar nuestro trabajo compartiendo en las redes sociales, haciendo una donación y comprando nuestro libro y curso electrónico.
Pytorch datos distribuidos en paralelo
El paralelo de datos distribuidos es multiproceso y funciona tanto para el entrenamiento de una sola máquina como para el de varias máquinas. en pytorch, nn.parallel.DistributedDataParallel paraleliza el módulo dividiendo la entrada entre los dispositivos especificados. Este módulo también es adecuado para el entrenamiento de varios nodos y varias GPU. Aquí, solo experimenté con un solo nodo (1 máquina con 4 GPU).
La principal diferencia aquí es que cada GPU es manejada por un proceso. Los parámetros nunca se transmiten entre procesos, solo gradientes.
El módulo se replica en cada máquina y cada dispositivo. Durante el paso hacia adelante, cada trabajador (GPU) procesa los datos y calcula su propio gradiente localmente. Durante el paso hacia atrás, se promedian los gradientes de cada nodo. Finalmente, cada trabajador realiza una actualización de parámetros y envía a todos los demás nodos la actualización de parámetros calculada.
El módulo realiza una todo-reducir pisar gradientes y asume que serán modificados por el optimizador en todos los procesos de la misma manera.
A continuación se encuentran las pautas para convertir su secuencia de comandos de GPU única en un entrenamiento de GPU múltiple.
Paso 1: inicializar los procesos de aprendizaje distribuido
def init_distributed():
dist_url = "env://"
rank = int(os.environ["RANK"])
world_size = int(os.environ['WORLD_SIZE'])
local_rank = int(os.environ['LOCAL_RANK'])
dist.init_process_group(
backend="nccl",
init_method=dist_url,
world_size=world_size,
rank=rank)
torch.cuda.set_device(local_rank)
dist.barrier()
Esta inicialización funciona cuando lanzamos nuestro script con torch.distributed.launch (Pytorch 1.7 y 1.8) o torch.run (Pytorch 1.9+) de cada nodo (aquí 1).
Paso 2: Envuelva el modelo usando DDP
net = torchvision.models.resnet50(False).cuda()
net = nn.SyncBatchNorm.convert_sync_batchnorm(net)
local_rank = int(os.environ['LOCAL_RANK'])
net = nn.parallel.DistributedDataParallel(net, device_ids=[local_rank])
Si cada proceso tiene el rango local correcto, tensor.cuda() o model.cuda() se puede llamar correctamente en todo el script.
Paso 3: Use un DistributedSampler en su DataLoader
import torch
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
import torch.nn as nn
def create_data_loader_cifar10():
transform = transforms.Compose(
[
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
batch_size = 256
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
train_sampler = DistributedSampler(dataset=trainset, shuffle=True)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
sampler=train_sampler, num_workers=10, pin_memory=True)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
test_sampler =DistributedSampler(dataset=testset, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, sampler=test_sampler, num_workers=10)
return trainloader, testloader
En modo distribuido, llamando al data_loader.sampler.set_epoch() método al comienzo de cada época antes creando el DataLoader iterator es necesario para hacer que el barajado funcione correctamente en varias épocas. En caso contrario, se utilizará siempre el mismo ordenamiento.
def train(net, trainloader):
print("Start training...")
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
epochs = 1
num_of_batches = len(trainloader)
for epoch in range(epochs):
trainloader.sampler.set_epoch(epoch)
En una forma más general:
for epoch in range(epochs):
data_loader.sampler.set_epoch(epoch)
train_one_epoch(...)
Buenas prácticas para DDP
Cualquier método que descargue datos debe estar aislado del proceso maestro. Cualquier método que realice E/S de archivos debe estar aislado del proceso maestro.
import torch.distributed as dist
import torch
def is_dist_avail_and_initialized():
if not dist.is_available():
return False
if not dist.is_initialized():
return False
return True
def save_on_master(*args, **kwargs):
if is_main_process():
torch.save(*args, **kwargs)
def get_rank():
if not is_dist_avail_and_initialized():
return 0
return dist.get_rank()
def is_main_process():
return get_rank() == 0
Según esta función, puede estar seguro de que algunos comandos solo se ejecutan desde el proceso principal:
if is_main_process():
Ejecutar script usando torch.distributed.launch o torch.run
$ python -m torch.distributed.launch --nproc_per_node=4 main_script.py
Se producirán errores. Asegúrese de eliminar cualquier proceso de entrenamiento distribuido no deseado al:
$ kill $(ps aux | grep main_script.py | grep -v grep | awk '')
Reemplazar main_script.py con su nombre de guión. Otra opción más sencilla es $ kill -9 PID. De lo contrario, puede ir a cosas más avanzadas, como matar todos los procesos relacionados con GPU CUDA cuando no se muestran en nvidia-smi
lsof /dev/nvidia* | awk '' | xargs -I kill
Esto es solo para el caso de que no pueda encontrar el PID del proceso que se ejecuta en la GPU.
Un muy buen libro sobre entrenamiento distribuido es Aprendizaje automático distribuido con Python: aceleración del entrenamiento y servicio de modelos con sistemas distribuidos por Guanhua Wang.
Entrenamiento de precisión mixta en Pytorch
La precisión mixta combina coma flotante (FP) 16 y FP 32 en diferentes pasos del entrenamiento. El entrenamiento FP16 también se conoce como entrenamiento de media precisión, que viene con un rendimiento inferior. La precisión mixta automática es literalmente lo mejor de ambos mundos: tiempo de entrenamiento reducido con un rendimiento comparable al FP32.
En Entrenamiento mixto de precisión, todas las operaciones computacionales (paso adelante, paso atrás, gradientes de peso) ven la versión casteada del FP16. Para hacerlo, es necesaria una copia del peso en FP32, así como calcular la pérdida en FP32 después del pase hacia adelante en FP16 para evitar overflows y underflows. Los gradientes de peso se proyectan de nuevo a FP32 para actualizar los pesos del modelo. Además, la pérdida en FP32 se escala para evitar el desbordamiento del gradiente antes de pasar a FP16 para realizar el pase hacia atrás. Como compensación, los pesos del FP32 se reducirán en el mismo escalar antes de la actualización del peso.
Estos son los cambios en la función de tren:
fp16_scaler = torch.cuda.amp.GradScaler(enabled=True)
for epoch in range(epochs):
trainloader.sampler.set_epoch(epoch)
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
images, labels = inputs.cuda(), labels.cuda()
optimizer.zero_grad()
with torch.cuda.amp.autocast():
outputs = net(images)
loss = criterion(outputs, labels)
fp16_scaler.scale(loss).backward()
fp16_scaler.step(optimizer)
fp16_scaler.update()
Resultados y Resumen
En un mundo paralelo utópico, N trabajadores darían una aceleración de N. Aquí puede ver que necesita 4 GPU en modo DistributedDataParallel para obtener una aceleración de 2X. El entrenamiento de precisión mixto normalmente proporciona una aceleración sustancial, pero la GPU A100 y otras arquitecturas de GPU basadas en Ampere tienen ganancias limitadas (por lo que he leído en línea).
Los resultados a continuación informan el tiempo en segundos para 1 época en CIFAR10 con un resnet50 (tamaño de lote 256, memoria GPU NVidia A100 de 40 GB):
| Tiempo en segundos | |
| GPU única (línea base) | 13.2 |
| GPU DataParallel 4 | 19.1 |
| DistributedDataParallel 2 GPU | 9.8 |
| DistributedDataParallel 4 GPU | 6.1 |
| DistributedDataParallel 4 GPU + precisión mixta | 6.5 |
Una nota muy importante aquí es que DistributedDataParallel utiliza un tamaño de lote efectivo de 4*256=1024 por lo que hace menos actualizaciones de modelos. Es por eso que creo que obtiene una precisión de validación mucho más baja (14 % en comparación con 27 % en la línea de base).
El código está disponible en GitHub si quieres jugar. Los resultados variarán según su hardware. Siempre se da el caso de que me perdí algo en mis experimentos. Si encuentra un defecto, hágamelo saber en nuestro servidor de discordia.
Estos hallazgos le brindarán un comienzo sólido para entrenar sus modelos. Espero que los encuentre útiles. Nos apoya compartiendo en las redes sociales, haciendo una donación, comprando nuestro libro o curso electrónico. Su ayuda nos ayudaría a producir más contenido gratuito y contenido de IA accesible. Como siempre, gracias por su interés en nuestro blog.
* Divulgación: tenga en cuenta que algunos de los enlaces anteriores pueden ser enlaces de afiliados y, sin costo adicional para usted, ganaremos una comisión si decide realizar una compra después de hacer clic.
Relacionado:
Precisión | Significado de precisión
Encuentro de realidad mixta de Zapbox: acercando la realidad mixta a todos
Maximización de la eficiencia y la precisión con la integración de NetSuite
