Entrenamiento con varias GPU con YOLOv5#

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 de Python>=3.8.0, incluido PyTorch>=1.8. Los modelos y los conjuntos de datos 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 desde el que comenzar el entrenamiento. Aquí seleccionamos 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 en 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)#

Antepon python -m torch.distributed.run --nproc_per_node al comando de entrenamiento y pasa después 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 utilizarán. En el ejemplo anterior, es 2.
• --batch es el tamaño total del lote, dividido equitativamente entre cada GPU. En el ejemplo anterior, es 64 / 2 = 32 por GPU.

El comando anterior utiliza las GPU 0...(N-1). Para controlar la visibilidad de los dispositivos mediante variables de entorno, establece CUDA_VISIBLE_DEVICES=2,3 (o cualquier otra lista) antes de iniciar.

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 utiliza ligeramente más 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 utilizando el mismo puerto. Especifica otro puerto con --master_port:

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

Resultados#

Resultados de la creación de perfiles de DDP en una instancia AWS EC2 P4d con 8 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 proporciona un escalado casi lineal en la velocidad de entrenamiento. Con 8 GPU, el entrenamiento finaliza 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 preinstalado con dependencias esenciales como CUDA, CUDNN, Python y PyTorch, para que puedas iniciar tus proyectos.

• Cuadernos de GPU gratuitos:
• 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 benchmarks. Garantizan un funcionamiento uniforme y fiable en macOS, Windows y Ubuntu, con pruebas que se realizan cada 24 horas y con cada nuevo commit.

Créditos#

Queremos dar las gracias a @MagicFrogSJTU, que hizo todo el trabajo pesado, y a @glenn-jocher por guiarnos durante el proceso.

Consulta también#

• Modo de entrenamiento - Aprende sobre el entrenamiento de 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 de Docker para el entrenamiento

Preguntas frecuentes#

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