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 bằng nhiều GPU trên một máy hoặc trên nhiều máy.

Trước khi bắt đầu#

Clone repository và cài đặt requirements.txt trong môi trường Python>=3.8.0, bao gồm PyTorch>=1.8. Model và dataset được tự động tải xuống từ bản phát hành YOLOv5 mới nhất.

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

Huấn luyện#

Chọn model pretrained để bắt đầu huấn luyện. Ở đây, chúng ta chọn YOLOv5s, một model nhỏ và nhanh. Xem bảng trong README để so sánh đầy đủ tất cả model. Chúng ta sẽ huấn luyện model này bằng Multi-GPU trên dataset COCO.

Một 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 chậm và hầu như không tăng tốc huấn luyện so với việc chỉ dùng một GPU.

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

Thêm python -m torch.distributed.run --nproc_per_node vào đầu lệnh huấn luyện, sau đó truyền các tham 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ố GPU cần sử dụng. Trong ví dụ trên, giá trị này là 2.
• --batch là tổng batch size, được chia đều cho từng GPU. Trong ví dụ trên, mỗi GPU nhận 64 / 2 = 32.

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

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ú#

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

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

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

• Nếu gặp RuntimeError: Address already in use, nguyên nhân thường là nhiều lần chạy huấn luyện đang sử dụng cùng một cổng. Hãy chỉ định cổng khác bằng --master_port:

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

Kết quả#

Kết quả profiling DDP trên instance AWS EC2 P4d với 8x A100 SXM4-40GB cho YOLOv5l trong 1 epoch 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

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

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 được cài đặt sẵn các dependency thiết yếu như CUDA, CUDNN, Python và PyTorch, giúp bạn bắt đầu dự án nhanh chóng.

• Notebook GPU miễn phí:
• Google Cloud: Hướng dẫn bắt đầu nhanh với GCP
• Amazon: Hướng dẫn bắt đầu nhanh với AWS
• Azure: Hướng dẫn bắt đầu nhanh với AzureML
• Docker: Hướng dẫn bắt đầu nhanh với Docker

Trạng thái dự án#

Huy hiệu này cho biết tất cả bài kiểm thử Continuous Integration (CI) của GitHub Actions cho YOLOv5 đều đã chạy 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 năng của YOLOv5 trên nhiều khía cạnh chính: huấn luyện, validation, inference, export và benchmark. Các bài kiểm thử đảm bảo hoạt động nhất quán và đáng tin cậy trên macOS, Windows và Ubuntu; chúng được thực hiện mỗi 24 giờ và sau mỗi commit mới.

Lời cảm ơn#

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

Xem thêm#

• Chế độ Train - Tìm hiểu cách huấn luyện model YOLO với Ultralytics
• Tinh chỉnh siêu tham số - Tối ưu hiệu suất model
• Hướng dẫn bắt đầu nhanh với Docker - Thiết lập môi trường Docker để huấn luyện

Câu hỏi thường gặp#

Hãy đọc danh sách kiểm tra bên dưới trước khi mở issue — việc này thường giúp tiết kiệm thời gian.