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

Обучение#

Выбери предварительно обученную модель, с которой начнёшь обучение. Здесь мы выбираем YOLOv5s — небольшую и быструю модель. Полное сравнение всех моделей см. в [таблице](https://ultralytics-translation-1.invalid нашего README. Мы будем обучать эту модель на датасете COCO с использованием нескольких GPU.

Одна GPU#

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

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

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

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

DataParallel работает медленно и почти не ускоряет обучение по сравнению с использованием одной GPU.

Режим Multi-GPU DistributedDataParallel (✅ рекомендуется)#

Добавь 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. В приведённом выше примере на каждую GPU приходится 64 / 2 = 32.

Приведённая выше команда использует 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 за помощь и руководство на этом пути.

См. также#

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

Часто задаваемые вопросы#

Прочитай приведённый ниже список перед созданием issue — это часто экономит время.