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çares

Clone o repositório e instale os requirements.txt em um ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Modelos e datasets são baixados automaticamente a partir da última release do YOLOv5.

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

Treinamento

Selecione um modelo pré-treinado para iniciar o treinamento. Aqui selecionamos o YOLOv5s, o modelo mais leve e rápido disponível. Consulte a tabela no nosso README para uma comparação completa de todos os modelos. Treinaremos este modelo com Multi-GPU no dataset COCO.

GPU Única

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

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

Passe múltiplos 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 treinamento em comparação ao uso de uma única GPU.

Modo Multi-GPU DistributedDataParallel (✅ recomendado)

Prefixe o comando de treinamento com python -m torch.distributed.run --nproc_per_node, então passe 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 entre cada GPU. No exemplo acima, isso é 64 / 2 = 32 por GPU.

O comando acima usa as GPUs 0...(N-1). Para controlar a visibilidade dos dispositivos através de variáveis de ambiente, defina 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 ao Windows não foi testado; Linux é recomendado.

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

• A GPU 0 utiliza um pouco mais de memória que as outras porque mantém o EMA e gerencia os checkpoints.

• Se você receber RuntimeError: Address already in use, geralmente significa que múltiplos treinamentos estão usando a mesma porta. Especifique uma porta diferente com --master_port:

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

Resultados

Resultados de perfilamento DDP em uma instância AWS EC2 P4d com 8x A100 SXM4-40GB para YOLOv5l por 1 epoch no COCO.

# prepare
t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --runtime=nvidia --ipc=host --gpus 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 mostrado nos resultados, usar DistributedDataParallel com múltiplas GPUs fornece um escalonamento quase linear na velocidade de treinamento. Com 8 GPUs, o treinamento é concluído aproximadamente 6,5 vezes mais rápido do que com uma única GPU, mantendo o mesmo uso de memória por dispositivo.

FAQ

Leia a lista de verificação abaixo antes de abrir uma issue — isso frequentemente economiza tempo.

Ambientes Suportados

A Ultralytics fornece uma variedade de ambientes prontos para uso, cada um pré-instalado com dependências essenciais como CUDA, CUDNN, Python e PyTorch, para dar início aos seus projetos.

• Notebooks com GPU Gratuita:
• Google Cloud: Guia de Início Rápido GCP
• Amazon: Guia de Início Rápido AWS
• Azure: Guia de Início Rápido AzureML
• Docker: Guia de Início Rápido Docker

Status do Projeto

Este selo indica que todos os testes de Integração Contínua (CI) do GitHub Actions do YOLOv5 estão passando com sucesso. Esses testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspectos fundamentais: treinamento, validação, inferência, exportação e benchmarks. Eles garantem uma operação consistente e confiá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 orientar ao longo do caminho.

Veja também

• Modo de Treinamento - Aprenda sobre o treinamento de modelos YOLO com a Ultralytics
• Ajuste de Hiperparâmetros - Otimize o desempenho do seu modelo
• Guia de Início Rápido Docker - Configure seu ambiente Docker para treinamento