如何使用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标志选择要使用的GPU，而其他标志仅确定我的深度神经网络的超参数。该脚本的作用是，对于for循环的每次迭代，它将在每个GPU上启动一个训练任务，并等待所有任务完成后再开始下一次迭代。问题在于，每个训练任务可能需要花费不同的时间来运行，因此，我同时使用少于8个GPU的时间将非常长，这会延长完成整个任务的时间。

我想知道是否有某种方法可以检测哪个GPU已完成其任务并在其上启动新任务，以便始终可以运行8个GPU。

非常感谢！

Answer 1

我看到您没有使用群集，这意味着GPU在您的本地计算机上。

在这种情况下，您可以使用此库：https://pypi.org/project/simple-gpu-scheduler/

希望这会有所帮助。