Huấn luyện Multi-GPU với YOLOv5

Hướng dẫn này giải thích cách huấn luyện YOLOv5 với nhiều GPU trên một máy duy nhất hoặc trên nhiều máy.

Trước khi Bạn Bắt đầu

Sao chép repo và cài đặt requirements.txt trong môi trường Python>=3.8.0, bao gồm PyTorch>=1.8. Models và datasets sẽ tự động tải xuống từ bản YOLOv5 release mới nhất.

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

Training

Chọn một model đã huấn luyện trước (pretrained) để bắt đầu huấn luyện. Ở đây chúng tôi chọn YOLOv5s, model nhỏ và nhanh nhất hiện có. Xem bảng trong README của chúng tôi để so sánh đầy đủ tất cả các model. Chúng tôi sẽ huấn luyện model này với Multi-GPU trên tập dữ liệu COCO.

Single GPU

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

Chế độ Multi-GPU DataParallel (⚠️ không khuyến nghị)

Truyền nhiều ID GPU vào --device để bật chế độ DataParallel:

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

DataParallel hoạt động chậm và hầu như không giúp tăng tốc độ huấn luyện so với việc sử dụng một GPU duy nhất.

Chế độ Multi-GPU DistributedDataParallel (✅ khuyến nghị)

Thêm tiền tố vào lệnh huấn luyện với python -m torch.distributed.run --nproc_per_node, sau đó truyền các đối số thông thường:

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 là số lượng GPU cần sử dụng. Trong ví dụ trên, số lượng là 2.
• --batch là tổng kích thước batch, được chia đều cho mỗi GPU. Trong ví dụ trên, đó là 64 / 2 = 32 cho mỗi GPU.

Lệnh trên sử dụng các GPU 0...(N-1). Để kiểm soát khả năng hiển thị của thiết bị thông qua biến môi trường thay vì cách trên, hãy thiết lập CUDA_VISIBLE_DEVICES=2,3 (hoặc bất kỳ danh sách nào khác) trước khi chạy lệnh.

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

Ghi chú

• Hỗ trợ Windows chưa được kiểm thử; khuyến nghị sử dụng Linux.

• --batch phải là bội số của số lượng GPU.

• GPU 0 sử dụng bộ nhớ nhiều hơn một chút so với các GPU khác vì nó duy trì EMA và xử lý việc lưu checkpoint.

• Nếu bạn gặp lỗi RuntimeError: Address already in use, thường có nghĩa là có nhiều tiến trình huấn luyện đang sử dụng cùng một cổng. Hãy chỉ định một cổng khác với --master_port:

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

Kết quả (Results)

Kết quả profiling DDP trên một AWS EC2 P4d instance với 8x A100 SXM4-40GB cho YOLOv5l trong 1 COCO epoch.

# 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

Như kết quả cho thấy, việc sử dụng DistributedDataParallel với nhiều GPU mang lại khả năng mở rộng gần như tuyến tính về tốc độ huấn luyện. Với 8 GPU, quá trình huấn luyện hoàn thành nhanh hơn khoảng 6,5 lần so với sử dụng một GPU duy nhất, trong khi vẫn duy trì mức sử dụng bộ nhớ trên mỗi thiết bị.

Câu hỏi thường gặp (FAQ)

Đọc danh sách kiểm tra dưới đây trước khi mở một Issue — nó thường giúp tiết kiệm thời gian.

Môi trường được hỗ trợ

Ultralytics cung cấp nhiều môi trường sẵn sàng sử dụng, mỗi môi trường đều được cài đặt sẵn các phụ thuộc thiết yếu như CUDA, CUDNN, Python, và PyTorch để giúp bạn bắt đầu dự án của mình.

• Notebook GPU miễn phí:
• Google Cloud: GCP Quickstart Guide
• Amazon: AWS Quickstart Guide
• Azure: AzureML Quickstart Guide
• Docker: Docker Quickstart Guide

Trạng thái dự án

Huy hiệu này cho biết tất cả các bài kiểm thử Tích hợp liên tục (CI) của YOLOv5 GitHub Actions đều đã vượt qua thành công. Các bài kiểm thử CI này kiểm tra nghiêm ngặt chức năng và hiệu suất của YOLOv5 trên nhiều khía cạnh quan trọng: huấn luyện, xác thực, suy luận, xuất mô hình, và điểm chuẩn. Chúng đảm bảo hoạt động nhất quán và tin cậy trên macOS, Windows và Ubuntu, với các bài kiểm thử được tiến hành mỗi 24 giờ và sau mỗi lần commit mới.

Ghi nhận (Credits)

Chúng tôi xin cảm ơn @MagicFrogSJTU, người đã thực hiện phần việc khó khăn nhất, và @glenn-jocher vì đã hướng dẫn chúng tôi trong suốt quá trình.

Xem thêm

• Train Mode - Tìm hiểu về cách huấn luyện các model YOLO với Ultralytics
• Hyperparameter Tuning - Tối ưu hóa hiệu suất của model
• Docker Quickstart Guide - Thiết lập môi trường Docker của bạn để huấn luyện