Entrenamiento de YOLOv5 con varias GPU#

Esta guía explica cómo entrenar YOLOv5 con varias GPU en una sola máquina o en varias máquinas.

Antes de empezar#

Clona el repositorio e instala requirements.txt en un entorno Python>=3.8.0 que incluya PyTorch>=1.8. Los modelos y los conjuntos de datos se descargan automáticamente desde la versión más reciente de YOLOv5.

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

Entrenamiento#

Elige un modelo preentrenado desde el que iniciar el entrenamiento. Aquí elegimos YOLOv5s, un modelo pequeño y rápido. Consulta la tabla de nuestro README para ver una comparación completa de todos los modelos. Entrenaremos este modelo con varias GPU usando el conjunto de datos COCO.

Una sola GPU#

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

Modo DataParallel con varias GPU (⚠️ no recomendado)#

Pasa varios 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 varias GPU (✅ recomendado)#

Añade python -m torch.distributed.run --nproc_per_node al principio del comando de entrenamiento y, a continuación, 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 que se van a usar. En el ejemplo anterior, es 2.
• --batch es el tamaño total del lote, que se divide equitativamente entre cada GPU. En el ejemplo anterior, corresponde a 64 / 2 = 32 por GPU.

El comando anterior usa las GPU 0...(N-1). Para controlar qué dispositivos están visibles mediante variables de entorno, establece CUDA_VISIBLE_DEVICES=2,3 (o cualquier otra lista) antes de iniciar 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#

• La compatibilidad con Windows no se ha probado; se recomienda Linux.

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

• La GPU 0 usa algo más de memoria que las demás porque mantiene la EMA y gestiona los puntos de control.

• Si aparece RuntimeError: Address already in use, normalmente significa que varias ejecuciones de entrenamiento están usando el mismo puerto. Especifica otro puerto con --master_port:

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

Resultados#

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

# prepare
t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --device nvidia.com/gpu=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 muestran los resultados, usar DistributedDataParallel con varias GPU permite escalar casi linealmente 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.

Entornos compatibles#

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

• Cuadernos gratuitos con GPU:
• 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 superan correctamente. Estas pruebas de CI comprueban rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: entrenamiento, validación, inferencia, exportación y evaluaciones comparativas. Garantizan un funcionamiento uniforme y fiable en macOS, Windows y Ubuntu, con pruebas cada 24 horas y tras cada nueva confirmación.

Agradecimientos#

Queremos dar las gracias a @MagicFrogSJTU, que se encargó de la mayor parte del trabajo, y a @glenn-jocher, que nos guio durante el proceso.

Consulta también#

• Modo de 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

Preguntas frecuentes#

Lee la siguiente lista de comprobación antes de abrir una incidencia: a menudo te ahorrará tiempo.