Entrenamiento con múltiples GPU con YOLOv5

Esta guía explica cómo entrenar YOLOv5 con múltiples GPU en una sola máquina o a través de varias máquinas.

Antes de empezar

Clona el repositorio e instala requirements.txt en un entorno Python>=3.8.0, incluyendo PyTorch>=1.8. Los modelos y datasets se descargan automáticamente desde la última versión de YOLOv5.

git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install

Entrenamiento

Selecciona un modelo preentrenado para comenzar a entrenar. Aquí seleccionamos YOLOv5s, el modelo más pequeño y rápido disponible. Consulta nuestra tabla del README para obtener una comparativa completa de todos los modelos. Entrenaremos este modelo con múltiples GPU usando el dataset COCO.

GPU única

python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0

Modo DataParallel con múltiples GPU (⚠️ no recomendado)

Pasa múltiples ID de GPU a --device para activar el modo DataParallel:

python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1

DataParallel es lento y apenas acelera el entrenamiento en comparación con el uso de una sola GPU.

Modo DistributedDataParallel con múltiples GPU (✅ recomendado)

Prefija el comando de entrenamiento con python -m torch.distributed.run --nproc_per_node, y luego pasa los argumentos habituales:

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1

• --nproc_per_node es el número de GPU a utilizar. En el ejemplo anterior, es 2.
• --batch es el tamaño total del lote, dividido uniformemente entre cada GPU. En el ejemplo anterior, eso es 64 / 2 = 32 por GPU.

El comando anterior utiliza las GPU 0...(N-1). Para controlar la visibilidad del dispositivo a través de variables de entorno, configura CUDA_VISIBLE_DEVICES=2,3 (o cualquier otra lista) antes de lanzar el proceso.

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --device 2,3

python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --sync-bn

# On master machine 0
python -m torch.distributed.run --nproc_per_node G --nnodes N --node_rank 0 --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights ''

# On machine R
python -m torch.distributed.run --nproc_per_node G --nnodes N --node_rank R --master_addr "192.168.1.1" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights ''

Notas

• El soporte para Windows no ha sido probado; se recomienda Linux.

• --batch debe ser un múltiplo del número de GPU.

• La GPU 0 utiliza un poco más de memoria que las demás porque mantiene el EMA y gestiona los puntos de control (checkpointing).

• Si obtienes RuntimeError: Address already in use, generalmente significa que varias ejecuciones de entrenamiento están usando el mismo puerto. Especifica un puerto diferente con --master_port:

  python -m torch.distributed.run --master_port 1234 --nproc_per_node 2 ...

Results

Resultados de perfilado DDP en una instancia AWS EC2 P4d con 8x A100 SXM4-40GB para YOLOv5l durante 1 época de COCO.

# prepare
t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --runtime=nvidia --ipc=host --gpus all -v "$(pwd)"/coco:/usr/src/coco $t
pip3 install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
cd .. && rm -rf app && git clone https://github.com/ultralytics/yolov5 -b master app && cd app
cp data/coco.yaml data/coco_profile.yaml

# profile
python train.py --batch-size 16 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0
python -m torch.distributed.run --nproc_per_node 2 train.py --batch-size 32 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1
python -m torch.distributed.run --nproc_per_node 4 train.py --batch-size 64 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1,2,3
python -m torch.distributed.run --nproc_per_node 8 train.py --batch-size 128 --data coco_profile.yaml --weights yolov5l.pt --epochs 1 --device 0,1,2,3,4,5,6,7

Como se muestra en los resultados, el uso de DistributedDataParallel con múltiples GPU proporciona una escala casi lineal en la velocidad de entrenamiento. Con 8 GPU, el entrenamiento se completa aproximadamente 6,5 veces más rápido que con una sola GPU, manteniendo el mismo uso de memoria por dispositivo.

Preguntas frecuentes

Lee la lista de verificación a continuación antes de abrir un issue; a menudo ahorra tiempo.

Entornos compatibles

Ultralytics ofrece una gama de entornos listos para usar, cada uno preinstalado con dependencias esenciales como CUDA, CUDNN, Python y PyTorch, para poner en marcha tus proyectos.

• Notebooks con GPU gratuita:
• Google Cloud: Guía de inicio rápido de GCP
• Amazon: Guía de inicio rápido de AWS
• Azure: Guía de inicio rápido de AzureML
• Docker: Guía de inicio rápido de Docker

Estado del proyecto

Esta insignia indica que todas las pruebas de Integración Continua (CI) de GitHub Actions de YOLOv5 se están superando con éxito. Estas pruebas de CI comprueban rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: entrenamiento, validación, inferencia, exportación y benchmarks. Garantizan un funcionamiento coherente y fiable en macOS, Windows y Ubuntu, con pruebas realizadas cada 24 horas y con cada nuevo commit.

Créditos

Nos gustaría agradecer a @MagicFrogSJTU, quien realizó todo el trabajo pesado, y a @glenn-jocher por guiarnos en el camino.

Ver también

• Modo entrenamiento - Aprende a entrenar modelos YOLO con Ultralytics
• Ajuste de hiperparámetros - Optimiza el rendimiento de tu modelo
• Guía de inicio rápido de Docker - Configura tu entorno Docker para el entrenamiento