Treinamento do YOLOv5 em várias GPUs#

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

Antes de começar#

Clone o repositório e instale requirements.txt em um ambiente Python>=3.8.0, incluindo PyTorch>=1.8. Os modelos e conjuntos de dados são baixados automaticamente da versão mais recente 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, um modelo pequeno e rápido. Consulte a tabela do nosso README para ver uma comparação completa de todos os modelos. Vamos treinar esse modelo com várias GPUs no conjunto de dados COCO.

GPU única#

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

Modo DataParallel com várias GPUs (⚠️ não recomendado)#

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

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

Adicione python -m torch.distributed.run --nproc_per_node como prefixo ao comando de treinamento e, em seguida, 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 usar. 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 usa as GPUs 0...(N-1). Para controlar quais dispositivos ficam visíveis por meio 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 ''

Observações#

• O suporte ao Windows não foi testado; recomendamos o Linux.

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

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

• Se você receber RuntimeError: Address already in use, geralmente isso significa que várias execuções de treinamento estão usando a mesma porta. Especifique outra porta 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 durante 1 época 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 mostram os resultados, usar DistributedDataParallel com várias GPUs proporciona um aumento quase linear na velocidade de treinamento. Com 8 GPUs, o treinamento termina aproximadamente 6,5 vezes mais rápido do que com uma única GPU, mantendo o mesmo uso de memória por dispositivo.

Ambientes compatíveis#

A Ultralytics oferece uma variedade de ambientes prontos para uso, cada um com dependências essenciais pré-instaladas, como CUDA, CUDNN, Python e PyTorch, para você começar seus projetos.

• Notebooks gratuitos com GPU:
• 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 emblema indica que todos os testes de Integração contínua (CI) das GitHub Actions da YOLOv5 estão a passar. Estes testes de CI verificam rigorosamente a funcionalidade e o desempenho da YOLOv5 em vários aspetos importantes: treino, validação, inferência, exportação e avaliação de desempenho. Os testes garantem um funcionamento consistente e fiável em macOS, Windows e Ubuntu e são realizados a cada 24 horas e sempre que é criado um novo commit.

Agradecimentos#

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

Veja também#

• Modo de treinamento - Saiba como treinar modelos YOLO com a Ultralytics
• Ajuste de hiperparâmetros - Otimize o desempenho do seu modelo
• Guia de início rápido do Docker - Configure seu ambiente Docker para treinamento

Perguntas frequentes#

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