Treino com várias GPUs usando YOLOv5#

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

Antes de Começares#

Clona o repositório e instala o requirements.txt num ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Os modelos e os conjuntos de dados são transferidos automaticamente a partir da versão mais recente do YOLOv5.

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

Treino#

Seleciona um modelo pré-treinado a partir do qual iniciar o treino. Aqui selecionamos YOLOv5s, um modelo pequeno e rápido. Consulta a tabela do nosso README para ver uma comparação completa de todos os modelos. Vamos treinar este modelo com várias GPUs no conjunto de dados COCO.

Uma GPU#

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

Modo DataParallel com várias GPUs (⚠️ 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

O DataParallel é lento e quase não acelera o treino em comparação com a utilização de uma única GPU.

Modo DistributedDataParallel com várias GPUs (✅ recomendado)#

Coloca python -m torch.distributed.run --nproc_per_node no início do comando de treino e 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 lote, dividido igualmente entre cada GPU. No exemplo acima, são 64 / 2 = 32 por GPU.

O comando acima utiliza as 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; recomenda-se Linux.

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

• A GPU 0 utiliza ligeiramente mais memória do que as outras porque mantém a EMA e gere os pontos de verificação.

• Se obtiveres RuntimeError: Address already in use, normalmente significa que várias 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 do perfil de DDP numa instância AWS EC2 P4d com 8x A100 SXM4-40GB para YOLOv5l durante 1 época de 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 mostram os resultados, utilizar DistributedDataParallel com várias GPUs proporciona uma escalabilidade quase linear na velocidade de treino. Com 8 GPUs, o treino termina aproximadamente 6,5 vezes mais depressa do que com uma única GPU, mantendo a mesma utilização de memória por dispositivo.

Ambientes suportados#

A Ultralytics disponibiliza vários ambientes prontos a utilizar, cada um com dependências essenciais pré-instaladas, como CUDA, CUDNN, Python e PyTorch, para iniciares os teus projetos.

• Notebooks GPU gratuitos:
• 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

Estado do projeto#

Este selo indica que todos os testes de Integração Contínua (CI) das 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 importantes: treino, validação, inferência, exportação e benchmarks. Garantem um funcionamento consistente e fiável no macOS, Windows e Ubuntu, com testes realizados a cada 24 horas e após cada novo commit.

Créditos#

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

Ver também#

• Modo de treino - Aprende sobre o treino de modelos YOLO com a Ultralytics
• Ajuste de hiperparâmetros - Otimiza o desempenho do teu modelo
• Guia de início rápido do Docker - Configura o teu ambiente Docker para o treino

Perguntas frequentes#

Lê a lista de verificação abaixo antes de abrir um problema — muitas vezes poupa tempo.