Treinamento multi-GPU com YOLOv5#

Este guia explica como treinar o YOLOv5 com múltiplas GPUs em uma única máquina ou em várias máquinas.

Antes de começar#

Clona o repositório e instala o requirements.txt num ambiente Python>=3.8.0, incluindo o PyTorch>=1.8. Os Models e os datasets são descarregados automaticamente a partir do mais recente release do YOLOv5.

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

Treinamento#

Seleciona um modelo pré-treinado para começar a treinar. Aqui selecionamos o YOLOv5s, um modelo pequeno e rápido. Vê a nossa table do README para uma comparação completa de todos os modelos. Vamos treinar este modelo com Multi-GPU no conjunto de dados COCO.

GPU única#

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

Modo Multi-GPU DataParallel (⚠️ não recomendado)#

Passa vários IDs de GPU para --device para ativar o modo DataParallel:

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

DataParallel é lento e quase não acelera o treinamento em comparação ao uso de uma única GPU.

Modo Multi-GPU DistributedDataParallel (✅ recomendado)#

Adiciona o prefixo python -m torch.distributed.run --nproc_per_node ao comando de treino e, em seguida, passa os argumentos habituais:

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 é o número de GPUs a utilizar. No exemplo acima, é 2.
• --batch é o tamanho total do batch, dividido igualmente por cada GPU. No exemplo acima, isso corresponde a 64 / 2 = 32 por GPU.

O comando acima utiliza os GPUs 0...(N-1). Para controlar a visibilidade dos dispositivos através de variáveis de ambiente, define CUDA_VISIBLE_DEVICES=2,3 (ou qualquer outra lista) antes de iniciar.

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

Notas#

• O suporte para Windows não foi testado; o Linux é recomendado.

• --batch deve ser um múltiplo do número de GPUs.

• A GPU 0 usa um pouco mais de memória do que as outras porque mantém o EMA e gerencia os check-points.

• Se obtiveres RuntimeError: Address already in use, significa habitualmente que múltiplas execuções de treino estão a utilizar a mesma porta. Especifica uma porta diferente com --master_port:

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

Resultados#

Resultados de profiling DDP numa AWS EC2 P4d instance com 8x A100 SXM4-40GB para o YOLOv5l durante 1 epoch do 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

Como demonstrado nos resultados, utilizar DistributedDataParallel com múltiplos GPUs proporciona um dimensionamento quase linear na velocidade de treino. Com 8 GPUs, o treino conclui-se aproximadamente 6,5 vezes mais rápido do que com um único GPU, mantendo a mesma utilização de memória por dispositivo.

Ambientes suportados#

A Ultralytics disponibiliza uma gama de ambientes prontos a usar, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python e PyTorch, para dar início aos teus projetos.

• Notebooks de GPU Grátis:
• Google Cloud: Guia de Início Rápido do GCP
• Amazon: Guia de Início Rápido da AWS
• Azure: Guia de Início Rápido do AzureML
• Docker: Guia de Início Rápido do Docker

Status do projeto#

Este crachá indica que todos os testes de Integração Contínua (CI) do GitHub Actions do YOLOv5 estão a passar com sucesso. Estes testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspetos fundamentais: treino, validação, inferência, exportação e testes de desempenho. Garantem um funcionamento consistente e fiável no macOS, Windows e Ubuntu, com testes realizados a cada 24 horas e a cada novo commit.

Créditos#

Gostaríamos de agradecer a @MagicFrogSJTU, que fez todo o trabalho pesado, e a @glenn-jocher por nos guiar ao longo do caminho.

Veja também#

• Train Mode - Aprende sobre o treino de modelos YOLO com o Ultralytics
• Hyperparameter Tuning - Otimiza o desempenho do teu modelo
• Docker Quickstart Guide - Configura o teu ambiente Docker para o treino

FAQ#

Leia o checklist abaixo antes de abrir uma Issue — isso geralmente economiza tempo.