本篇記錄如何使用多張 GPU 顯示卡,加速 TensorFlow Object Detection API 模型訓練的過程。

雖然 TensorFlow Object Detection API 已經有支援多張 GPU 卡平行計算的功能,但是缺乏說明文件,所以我自己也不是非常確定該怎麼用,以下只是我目前嘗試出來的方式,僅供參考。


這裡我們接續之前的 TensorFlow Object Detection API 自行訓練模型教學,將 Oxford-IIIT Pet Dataset 的範例改成多 GPU 的版本。

使用多 GPU 卡時,會把每個 batch 的資料分散至每張 GPU 卡,也就是可以讓 batch size 變大,所以在 train_configbatch_size 要自己修改一下,要讓資料可以平均分散至每張 GPU 卡,最簡單的修改方式就是看自己有多少張 GPU 卡,就乘以多少。

在預設的設定值中,batch_size 的值是 1,現在我打算用 3 張 GPU 卡,所以就改成 3

train_config: {
  batch_size: 3
  optimizer {
    # [略]
  }
  # [略]
}

由於 batch size 變大了,收斂的速度應該也會變快,所以 optimizer 內的參數應該也是要改的,不過要怎麼改就要看實際情況而定,這個部分就自己看著辦。

改好設定檔之後,接著就可以使用多張 GPU 卡平行運算,使用的指令幾乎沒有變,只是加上兩個參數:

  • --num_clones:指定 GPU 卡的數量。
  • --ps_tasks:指定參數伺服器的數量。

另外再以 CUDA_VISIBLE_DEVICES 指定要使用哪幾張 GPU 卡,完整的指令搞如下:

# 設定檔路徑
PIPELINE_CONFIG="object_detection/data/faster_rcnn_resnet101_pets.config"

# 訓練結果放置路徑
MY_MODEL_DIR="my_model"

# 使用前三張 GPU 卡進行訓練
CUDA_VISIBLE_DEVICES=0,1,2 python object_detection/train.py \
  --logtostderr \
  --pipeline_config_path=${PIPELINE_CONFIG} \
  --train_dir=${MY_MODEL_DIR}/train \
  --num_clones=3 --ps_tasks=1

# 使用第四張 GPU 卡進行驗證
CUDA_VISIBLE_DEVICES=3 python object_detection/eval.py \
  --logtostderr \
  --pipeline_config_path=${PIPELINE_CONFIG} \
  --checkpoint_dir=${MY_MODEL_DIR}/train \
  --eval_dir=${MY_MODEL_DIR}/eval

指令執行之後,我們可以使用 nvidia-smi 來查看每一張 GPU 卡的使用情況,看看是不是真的有同時用到三張 GPU 卡訓練:

nvidia-smi
Tue Dec 26 15:35:34 2017       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  TITAN Xp            Off  | 00000000:02:00.0 Off |                  N/A |
| 44%   72C    P2   186W / 250W |  11763MiB / 12189MiB |     84%      Default |
+-------------------------------+----------------------+----------------------+
|   1  TITAN Xp            Off  | 00000000:03:00.0 Off |                  N/A |
| 47%   76C    P2   124W / 250W |  11763MiB / 12189MiB |     76%      Default |
+-------------------------------+----------------------+----------------------+
|   2  TITAN Xp            Off  | 00000000:83:00.0 Off |                  N/A |
| 42%   70C    P2   131W / 250W |  11763MiB / 12189MiB |     72%      Default |
+-------------------------------+----------------------+----------------------+
|   3  TITAN Xp            Off  | 00000000:84:00.0 Off |                  N/A |
| 23%   41C    P2    72W / 250W |  11761MiB / 12189MiB |     51%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0     22192      C   python                                     11751MiB |
|    1     22192      C   python                                     11751MiB |
|    2     22192      C   python                                     11751MiB |
|    3     29266      C   python                                     11751MiB |
+-----------------------------------------------------------------------------+

在訓練的過程中,如果仔細觀察每步的計算速度,應該會發現使用一張 GPU 卡的速度會跟使用多張 GPU 卡差不多。這是單張 GPU 卡的速度:

單張 GPU 卡計算速度

這是同時使用三張 GPU 卡的計算速度:

三張 GPU 卡計算速度

速度會差不多是正常的(理論上多張 GPU 卡運算的速度會變慢一點點),因為兩邊的 batch size 不同,若每步的計算時間差不多,就表示實際的計算速度有增加。

參考資料:GitHub