Treinamento com Múltiplas GPUs com YOLOv5

Este guia explica como usar corretamente múltiplas GPUs para treinar um conjunto de dados com YOLOv5 🚀 em uma ou mais máquinas.

Antes de começar

Clone o repositório e instale requirements.txt em um Python>=3.8.0 ambiente, incluindo PyTorch>=1.8. Modelos e datasets download automaticamente a partir do YOLOv5 mais recente release.

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

Treinamento

Selecione um modelo pré-treinado para começar o treinamento. Aqui selecionamos YOLOv5s, o modelo mais leve e rápido disponível. Veja nossa tabela no README para uma comparação completa de todos os modelos. Treinaremos este modelo com múltiplas GPUs no COCO.

GPU Única

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

Modo Múltiplas GPUs DataParallel (⚠️ não recomendado)

Você pode aumentar o device para usar múltiplas GPUs no modo DataParallel.

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

Este método é lento e dificilmente acelera o treinamento em comparação ao uso de apenas 1 GPU.

Modo Múltiplas GPUs Modo DistributedDataParallel (✅ recomendado)

Você terá que passar python -m torch.distributed.run --nproc_per_node, seguido pelos argumentos usuais.

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 especifica quantas GPUs você gostaria de usar. No exemplo acima, são 2.
• --batch é o batch-size total. Ele será dividido igualmente para cada GPU. No exemplo acima, é 64/2=32 por GPU.

O código acima usará as GPUs 0... (N-1). Você também pode definir CUDA_VISIBLE_DEVICES=2,3 (ou qualquer outra lista) antes de executar o comando se preferir controlar a visibilidade do dispositivo através de variáveis de ambiente.

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 é testado, o Linux é recomendado.

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

• A GPU 0 consumirá um pouco mais de memória do que as outras GPUs, pois mantém o EMA e é responsável por criar checkpoints, etc.

• Se você receber um RuntimeError: Address already in use, pode ser porque você está executando vários treinamentos ao mesmo tempo. Para corrigir isso, simplesmente use um número de porta diferente adicionando --master_port como abaixo:

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

Results

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

# 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 Modo DistributedDataParallel com múltiplas GPUs proporciona um dimensionamento 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

Se ocorrer um erro, por favor, leia a lista de verificação abaixo primeiro! (Pode economizar seu tempo)

Ambientes Suportados

A Ultralytics fornece uma gama de ambientes prontos para uso, cada um pré-instalado com dependências essenciais, como CUDA, CUDNN, Python, e Modelos , para impulsionar os teus 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 YOLOv5 GitHub Actions foram concluídos com sucesso. Esses testes de CI verificam rigorosamente a funcionalidade e o desempenho do YOLOv5 em vários aspectos fundamentais: training, validação, inferência, export, e benchmarks. Eles garantem uma operação consistente e confiável no macOS, Windows e Ubuntu, com testes conduzidos 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

• Modo de Treino (Train Mode) - 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 Docker - Configura o teu ambiente Docker para treino