Entrenamiento con varias GPU con YOLOv5#

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

Antes de empezar#

Clona el repositorio e instala requirements.txt en un entorno de Python>=3.8.0, incluyendo PyTorch>=1.8. Los modelos y conjuntos de datos se descargan automáticamente desde el último lanzamiento de YOLOv5.

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

Entrenamiento#

Selecciona un modelo preentrenado con el que empezar a entrenar. Aquí seleccionamos YOLOv5s, un modelo pequeño y rápido. Consulta la tabla de nuestro archivo README para ver una comparación completa de todos los modelos. Entrenaremos este modelo con múltiples 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 múltiples GPU (⚠️ no recomendado)#

Pasa múltiples ID de GPU a --device para habilitar 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)#

Antepone python -m torch.distributed.run --nproc_per_node al comando de entrenamiento 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 que se van a usar. En el ejemplo anterior, es 2.
• --batch es el tamaño total del lote, dividido equitativamente 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 de los dispositivos mediante variables de entorno en su lugar, configura 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#

• El soporte para Windows no está 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 la EMA y gestiona los puntos de control (checkpoints).

• Si obtienes RuntimeError: Address already in use, normalmente significa que múltiples 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 ...

Resultados#

Resultados de perfiles 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 --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 se muestra en los resultados, el uso de DistributedDataParallel con múltiples GPU proporciona un escalado 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 al mismo tiempo el mismo uso de memoria por dispositivo.

Entornos admitidos#

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: AWS Quickstart Guide
• Azure: AzureML Quickstart Guide
• Docker: Guía de inicio rápido de Docker

Estado del proyecto#

Este distintivo indica que todas las pruebas de Integración Continua (CI) de YOLOv5 GitHub Actions se están superando con éxito. Estas pruebas de CI comprueban rigurosamente la funcionalidad y el rendimiento de YOLOv5 en varios aspectos clave: training, validation, inference, export 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 hizo todo el trabajo pesado, y a @glenn-jocher por guiarnos en el camino.

Ver 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 de Docker para el entrenamiento

FAQ#

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