TensorFlow Object Detection API 多 GPU 卡平行計算，加速模型訓練速度教學

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

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

使用多 GPU 卡時，會把每個 batch 的資料分散至每張 GPU 卡，也就是可以讓 batch size 變大，所以在 train_config 的 batch_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 卡運算的速度會變慢一點點），因為兩邊的 batch size 不同，若每步的計算時間差不多，就表示實際的計算速度有增加。

參考資料：GitHub