如何使用 Bash 腳本保持所有 GPU 設備運行任務？

Question

我有一個 8-GPU 服務器，我想同時在每個服務器上訓練一個神經網絡。 我有幾十個這樣的網絡要訓練，我想安排訓練任務。 目前我正在為此調度任務編寫自己的 bash 腳本。

for l1 in {1e-4,2e-4,5e-4,1e-3}; do

      python train.py --lr $l1 --attr 0 --device 0 &
      python train.py --lr $l1 --attr 1 --device 1 &
      python train.py --lr $l1 --attr 2 --device 2 &
      python train.py --lr $l1 --attr 3 --device 3 &
      python train.py --lr $l1 --attr 4 --device 4 &
      python train.py --lr $l1 --attr 5 --device 5 &
      python train.py --lr $l1 --attr 6 --device 6 &
      python train.py --lr $l1 --attr 7 --device 7

      sleep 1
      wait 
done 

在上面的腳本中，--device 標志選擇--device使用，而其他標志只是確定我的深度神經網絡的超參數。 該腳本的作用是，對於 for 循環的每次迭代，它會在每個 GPU 上啟動一個訓練任務，並等待所有這些任務完成，然后再開始下一次迭代。 問題是，每個訓練任務可能需要不同的時間來運行，因此我會同時使用少於 8 個 GPU 的大量時間，這會延長整個任務完成的時間。

我想知道是否有某種方法可以讓我檢測到哪個 GPU 已完成其任務並在其上啟動一個新任務，這樣我就可以始終運行 8 個 GPU。

非常感謝！

Answer 1

我看到您沒有使用集群，這意味着 GPU 在您的本地計算機上。

在這種情況下，您可以使用這個庫： https://pypi.org/project/simple-gpu-scheduler/

希望這可以幫助。

Answer 2

剛剛發現 Ray 是一款出色的 package，可用於管理您的實驗。 （ https://github.com/ray-project/ray ）