Обучение с несколькими 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 # install

Training#

Выбери предварительно обученную модель для начала обучения. Здесь мы выбираем YOLOv5s, маленькую и быструю модель. Полное сравнение всех моделей см. в нашей таблице в README. Мы будем обучать эту модель на нескольких GPU на наборе данных COCO.

Один GPU#

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

Режим DataParallel для нескольких GPU (⚠️ не рекомендуется)#

Передай несколько идентификаторов GPU в --device, чтобы включить режим DataParallel:

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

DataParallel работает медленно и почти не ускоряет обучение по сравнению с использованием одного 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 (или любой другой список) перед запуском.

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

Примечания#

• Поддержка 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 ...

Результаты#

Результаты профилирования DDP на инстансе AWS EC2 P4d с 8x A100 SXM4-40GB для модели YOLOv5l за 1 эпоху 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

Как показано в результатах, использование DistributedDataParallel с несколькими GPU обеспечивает практически линейное масштабирование скорости обучения. С 8 GPU обучение завершается примерно в 6.5 раз быстрее, чем с одним GPU, при сохранении прежнего объема использования памяти на устройство.

Поддерживаемые окружения#

Ultralytics предоставляет множество готовых к использованию сред, в каждую из которых предварительно установлены основные зависимости, такие как CUDA, CUDNN, Python и PyTorch, чтобы помочь тебе быстро начать проекты.

• Бесплатные GPU ноутбуки:
• Google Cloud: Краткое руководство по GCP
• Amazon: Краткое руководство по AWS
• Azure: Краткое руководство по AzureML
• Docker: Краткое руководство по Docker

Статус проекта#

Этот значок указывает на то, что все тесты непрерывной интеграции (CI) GitHub Actions для YOLOv5 успешно пройдены. Эти тесты CI тщательно проверяют функциональность и производительность YOLOv5 по различным ключевым аспектам: обучение, валидация, инференс, экспорт и бенчмарки. Они обеспечивают стабильную и надежную работу на macOS, Windows и Ubuntu, причем тесты проводятся каждые 24 часа и при каждом новом коммите.

Авторы#

Мы хотели бы поблагодарить @MagicFrogSJTU, который проделал всю основную работу, и @glenn-jocher за руководство на каждом этапе.

Смотри также#

• Режим обучения — узнай об обучении моделей YOLO с помощью Ultralytics
• Настройка гиперпараметров — оптимизируй производительность своей модели
• Руководство по быстрому старту Docker — настрой среду Docker для обучения

FAQ#

Ознакомься с чек-листом ниже перед созданием тикета (Issue) — это часто экономит время.