Entraînement multi-GPU avec YOLOv5#

Ce guide explique comment entraîner YOLOv5 avec plusieurs GPU sur une seule machine ou sur plusieurs machines.

Avant de commencer#

Clone le dépôt et installe requirements.txt dans un environnement Python>=3.8.0, avec PyTorch>=1.8. Les modèles et les jeux de données sont téléchargés automatiquement depuis la dernière version de YOLOv5.

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

Entraînement#

Choisis un modèle préentraîné pour commencer l'entraînement. Nous choisissons ici YOLOv5s, un modèle petit et rapide. Consulte le tableau de notre README pour comparer tous les modèles. Nous entraînerons ce modèle sur le jeu de données COCO avec plusieurs GPU.

Un seul GPU#

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

Mode DataParallel multi-GPU (⚠️ non recommandé)#

Transmets plusieurs identifiants de GPU à --device pour activer le mode DataParallel :

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

DataParallel est lent et accélère à peine l'entraînement par rapport à l'utilisation d'un seul GPU.

Mode DistributedDataParallel multi-GPU (✅ recommandé)#

Fais précéder la commande d'entraînement de python -m torch.distributed.run --nproc_per_node, puis transmets les arguments habituels :

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 correspond au nombre de GPU à utiliser. Dans l'exemple ci-dessus, il s'agit de 2.
• --batch correspond à la taille totale du lot, répartie équitablement entre chaque GPU. Dans l'exemple ci-dessus, cela représente 64 / 2 = 32 par GPU.

La commande ci-dessus utilise les GPU 0...(N-1). Pour contrôler la visibilité des périphériques à l'aide de variables d'environnement, définis plutôt CUDA_VISIBLE_DEVICES=2,3 (ou toute autre liste) avant le lancement.

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 ''

Remarques#

• La prise en charge de Windows n’a pas été testée ; Linux est recommandé.

• --batch doit être un multiple du nombre de GPU.

• Le GPU 0 utilise légèrement plus de mémoire que les autres, car il conserve l’EMA et gère les points de contrôle.

• Si tu obtiens RuntimeError: Address already in use, cela signifie généralement que plusieurs exécutions d’entraînement utilisent le même port. Indique un autre port avec --master_port :

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

Résultats#

Résultats du profilage DDP sur une instance AWS EC2 P4d avec 8x A100 SXM4-40GB pour YOLOv5l pendant 1 epoch 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

Comme le montrent les résultats, l’utilisation de DistributedDataParallel avec plusieurs GPU permet d’obtenir une accélération de l’entraînement presque linéaire. Avec 8 GPU, l’entraînement se termine environ 6,5 fois plus vite qu’avec un seul GPU, tout en conservant la même utilisation mémoire par appareil.

Environnements pris en charge#

Ultralytics propose un ensemble d’environnements prêts à l’emploi, chacun préinstallé avec les dépendances essentielles comme CUDA, CUDNN, Python et PyTorch, pour te permettre de démarrer tes projets.

• Notebooks GPU gratuits :
• Google Cloud : Guide de démarrage rapide GCP
• Amazon : Guide de démarrage rapide AWS
• Azure : Guide de démarrage rapide AzureML
• Docker : Guide de démarrage rapide Docker

État du projet#

Ce badge indique que tous les tests d’intégration continue (CI) de GitHub Actions pour YOLOv5 réussissent. Ces tests CI vérifient rigoureusement le fonctionnement et les performances de YOLOv5 dans plusieurs domaines clés : entraînement, validation, inférence, exportation et tests de performance. Ils garantissent un fonctionnement cohérent et fiable sur macOS, Windows et Ubuntu, grâce à des tests effectués toutes les 24 heures et à chaque nouveau commit.

Remerciements#

Nous remercions @MagicFrogSJTU, qui a fait tout le gros du travail, ainsi que @glenn-jocher, qui nous a guidés tout au long du parcours.

Voir aussi#

• Mode Entraînement - Découvre comment entraîner des modèles YOLO avec Ultralytics
• Réglage des hyperparamètres - Optimise les performances de ton modèle
• Guide de démarrage rapide Docker - Configure ton environnement Docker pour l’entraînement

FAQ#

Lis la liste de vérification ci-dessous avant d’ouvrir un ticket — elle permet souvent de gagner du temps.