YOLOv5によるマルチGPUトレーニング#

このガイドでは、1台のマシンまたは複数のマシンにまたがって複数のGPUを使用し、YOLOv5をトレーニングする方法を説明します。

開始する前に#

リポジトリをクローンし、Python>=3.8.0環境（PyTorch>=1.8を含む）でrequirements.txtをインストールします。モデルとデータセットは、最新のYOLOv5リリースから自動的にダウンロードされます。

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

トレーニング#

トレーニングを開始するための事前学習済みモデルを選択します。ここでは、小型で高速なモデルであるYOLOv5sを選択します。すべてのモデルを詳しく比較するには、READMEの表をご覧ください。このモデルをCOCOデータセットでマルチGPUトレーニングします。

シングルGPU#

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

マルチGPUのDataParallelモード（⚠️ 推奨されません）#

複数のGPU IDを--deviceに渡すと、DataParallelモードが有効になります。

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

DataParallelは低速で、シングルGPUの使用と比べてトレーニングをほとんど高速化できません。

マルチGPUのDistributedDataParallelモード（✅ 推奨）#

トレーニングコマンドの先頭にpython -m torch.distributed.run --nproc_per_nodeを付け、通常の引数を渡します。

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は使用するGPUの数です。上記の例では2です。
• --batchは合計バッチサイズで、各GPUに均等に分割されます。上記の例では、1つのGPUあたり64 / 2 = 32です。

上記のコマンドではGPU 0...(N-1)を使用します。代わりに環境変数でデバイスの表示を制御するには、起動前にCUDA_VISIBLE_DEVICES=2,3（またはほかのリスト）を設定します。

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

備考#

• Windowsのサポートは未検証のため、Linuxを推奨します。

• --batch はGPU数の倍数である必要があります。

• GPU 0はEMAを維持し、チェックポイント処理を行うため、ほかのGPUよりメモリをわずかに多く使用します。

• RuntimeError: Address already in use が発生した場合、通常は複数のトレーニング実行で同じポートを使用しています。--master_port を使って別のポートを指定してください：

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

結果#

YOLOv5lでCOCOの1 エポックを実行した際の、8基のA100 SXM4-40GBを搭載したAWS EC2 P4dインスタンス上でのDDPプロファイリング結果。

# 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

結果が示すように、複数のGPUでDistributedDataParallelを使用すると、トレーニング速度はほぼ線形にスケールします。8基のGPUを使用すると、デバイスあたりのメモリ使用量を維持したまま、単一GPUの場合と比べて約6.5倍速くトレーニングが完了します。

サポート対象の環境#

Ultralyticsでは、プロジェクトをすぐに開始できるよう、CUDA、CUDNN、Python、PyTorchなどの必須依存関係をあらかじめインストールした、すぐに使えるさまざまな環境を提供しています。

• 無料のGPUノートブック:
• Google Cloud: GCPクイックスタートガイド
• Amazon: AWSクイックスタートガイド
• Azure: AzureMLクイックスタートガイド
• Docker: Dockerクイックスタートガイド

プロジェクトのステータス#

このバッジは、すべてのYOLOv5 GitHub Actions継続的インテグレーション（CI）テストが正常に合格していることを示します。これらのCIテストでは、トレーニング、検証、推論、エクスポート、ベンチマークなど、さまざまな重要な側面におけるYOLOv5の機能とパフォーマンスを厳密に確認します。これにより、macOS、Windows、Ubuntuで一貫して信頼性の高い動作が保証されます。テストは24時間ごと、および新しいコミットごとに実施されます。

謝辞#

多くの作業を担ってくださった@MagicFrogSJTUと、取り組みを導いてくださった@glenn-jocherに感謝いたします。

関連項目#

• トレーニングモード - UltralyticsでYOLOモデルをトレーニングする方法について説明します
• ハイパーパラメータチューニング - モデルのパフォーマンスを最適化します
• Dockerクイックスタートガイド - トレーニング用のDocker環境をセットアップします

よくある質問#

Issueを作成する前に、以下のチェックリストを確認してください。時間の節約につながることがよくあります。