Обучение YOLOv5 на нескольких GPU#

В этом руководстве объясняется, как обучать 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 — небольшую и быструю модель. Полное сравнение всех моделей см. в таблице в 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 (epoch).

# 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

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

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

Благодарности#

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

См. также#

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

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

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