Обучение с несколькими GPU в YOLOv5
В этом руководстве объясняется, как обучать YOLOv5 с использованием нескольких GPU на одной машине или на нескольких машинах.
Перед началом
Клонируй репозиторий и установи requirements.txt в среде Python>=3.8.0, включая PyTorch>=1.8. Модели и датасеты загружаются автоматически из последнего релиза YOLOv5.
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # installDocker-образ Ultralytics рекомендуется для всех запусков обучения на нескольких GPU. См. Руководство по быстрому старту Docker.
torch.distributed.run заменяет torch.distributed.launch в PyTorch >= 1.9. Подробности см. в документации по распределенным вычислениям PyTorch.
Обучение
Выбери предобученную модель для начала обучения. Здесь мы выбираем YOLOv5s — самую маленькую и быструю из доступных моделей. Полное сравнение всех моделей смотри в таблице в нашем README. Мы будем обучать эту модель на нескольких GPU, используя датасет COCO.

Один GPU
python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0Режим DataParallel для нескольких GPU (⚠️ не рекомендуется)
Передай несколько ID GPU в --device, чтобы включить режим DataParallel:
python train.py --batch 64 --data coco.yaml --weights yolov5s.pt --device 0,1DataParallel работает медленно и почти не ускоряет обучение по сравнению с использованием одного GPU.
Режим DistributedDataParallel для нескольких GPU (✅ рекомендуется)
Добавь перед командой обучения префикс python -m torch.distributed.run --nproc_per_node, затем передай обычные аргументы:
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— это количество GPU, которые нужно использовать. В примере выше оно равно2.--batch— это общий размер батча, поровну разделенный на каждый GPU. В примере выше это64 / 2 = 32на каждый GPU.
Команда выше использует GPU 0...(N-1). Чтобы управлять видимостью устройств через переменные окружения, установи CUDA_VISIBLE_DEVICES=2,3 (или любой другой список) перед запуском.
Use specific GPUs (click to expand)
Передай --device и укажи конкретные ID GPU. В примере ниже используются GPU 2,3.
python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --device 2,3Use SyncBatchNorm (click to expand)
SyncBatchNorm может повысить точность при обучении на нескольких GPU, но значительно замедляет процесс. Этот режим доступен только при обучении через DistributedDataParallel на нескольких GPU.
Лучше всего использовать, когда размер батча на каждом GPU невелик (<= 8).
Чтобы включить SyncBatchNorm, передай аргумент --sync-bn:
python -m torch.distributed.run --nproc_per_node 2 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights '' --sync-bnUse Multiple machines (click to expand)
Это доступно только при обучении через DistributedDataParallel на нескольких GPU.
Перед продолжением убедись, что датасет, кодовая база и любые другие зависимости идентичны на всех машинах, а затем проверь, что машины могут видеть друг друга в сети.
Выбери главную машину (к которой будут подключаться остальные), запиши ее адрес (master_addr) и выбери порт (master_port). В примере ниже используются master_addr = 192.168.1.1 и master_port = 1234.
Затем выполни:
# 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 ''где G — количество GPU на каждой машине, N — количество машин, а R — порядковый номер (ранг) машины в диапазоне 0...(N-1). Например, при двух машинах с двумя GPU на каждой, установи G = 2, N = 2 и R = 1 на второй машине.
Обучение не начнется, пока все N машин не будут подключены. Вывод отображается только на главной машине.
Примечания
-
Поддержка Windows не тестировалась; рекомендуется Linux.
-
--batchдолжен быть кратен количеству GPU. -
GPU 0 использует немного больше памяти, чем остальные, так как он поддерживает EMA и обрабатывает создание чекпоинтов.
-
Если ты получаешь ошибку
RuntimeError: Address already in use, это обычно означает, что несколько процессов обучения используют один и тот же порт. Укажи другой порт с помощью--master_port:python -m torch.distributed.run --master_port 1234 --nproc_per_node 2 ...
Results
Результаты профилирования DDP на инстансе AWS EC2 P4d с 8x A100 SXM4-40GB для YOLOv5l за 1 эпоху на COCO.
Profiling code
# 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| GPUs A100 | batch-size | CUDA_mem устр0 (G) | COCO train | COCO val |
|---|---|---|---|---|
| 1x | 16 | 26GB | 20:39 | 0:55 |
| 2x | 32 | 26GB | 11:43 | 0:57 |
| 4x | 64 | 26GB | 5:57 | 0:55 |
| 8x | 128 | 26GB | 3:09 | 0:57 |
Как показывают результаты, использование DistributedDataParallel с несколькими GPU обеспечивает почти линейное масштабирование скорости обучения. С 8 GPU обучение проходит примерно в 6,5 раза быстрее, чем на одном, сохраняя при этом тот же объем использования памяти на устройство.
Часто задаваемые вопросы (FAQ)
Проверь список ниже, прежде чем открывать issue — это часто экономит время.
Checklist (click to expand)
- Ты прочитал это руководство целиком?
- Ты переклонировал кодовую базу? Код меняется ежедневно.
- Ты искал сообщение об ошибке? Возможно, кто-то уже столкнулся с такой же проблемой и поделился решением.
- Ты установил все зависимости (включая правильные версии Python и PyTorch)?
- Ты пробовал одну из поддерживаемых сред, перечисленных ниже?
- Ты пробовал использовать меньший датасет, например
coco128илиcoco2017, чтобы изолировать причину проблемы?
Если все вышеперечисленное в порядке, открой Issue с максимальным количеством подробностей, следуя шаблону.
Поддерживаемые среды
Ultralytics предоставляет ряд готовых к использованию сред, в которых уже предустановлены основные зависимости, такие как CUDA, CUDNN, Python и PyTorch, чтобы ты мог быстрее начать свои проекты.
- Бесплатные GPU ноутбуки:
- Google Cloud: Руководство по быстрому старту GCP
- Amazon: Руководство по быстрому старту AWS
- Azure: Руководство по быстрому старту AzureML
- Docker: Руководство по быстрому старту Docker
Статус проекта
Этот значок означает, что все тесты непрерывной интеграции (CI) YOLOv5 GitHub Actions успешно проходят. Эти CI-тесты строго проверяют функциональность и производительность YOLOv5 по различным ключевым аспектам: обучение, валидация, инференс, экспорт и бенчмарки. Они обеспечивают стабильную и надежную работу на macOS, Windows и Ubuntu, при этом тесты проводятся каждые 24 часа и при каждом новом коммите.
Авторы
Мы хотели бы поблагодарить @MagicFrogSJTU, который проделал основную тяжелую работу, и @glenn-jocher за руководство на всем пути.
Смотри также
- Режим обучения - Узнай об обучении моделей YOLO с помощью Ultralytics
- Настройка гиперпараметров - Оптимизируй производительность своей модели
- Руководство по быстрому старту Docker - Настрой среду Docker для обучения