Автоматическое распределение устройств CUDA для Tensorflow

Question

У меня есть несколько идентичных устройств CUDA на одном компьютере. На этом компьютере я запускаю несколько обучающих экземпляров tenorflow, каждый из которых использует одно и только одно устройство CUDA. Я хотел бы выделить одно и только одно устройство CUDA для экземпляра тензорного потока автоматически, в зависимости от того, что является бесплатным. Я хотел бы сделать это, когда запускается обучающий экземпляр tenorflow. В настоящее время я делаю это вручную, проверяя доступность и изменяя код в начале каждого обучающего экземпляра tenorflow:

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

Как я могу автоматически назначить первое (любое) неиспользуемое устройство CUDA для обучающего экземпляра тензорного потока? Значение в приведенном выше коде "0" должно быть изменено на "1", "2" et c. во время выполнения.

Answer 1

Все, что вам нужно, это планировщик графического процессора, вот один из них: https://pypi.org/project/simple-gpu-scheduler/

Я предполагаю, что вы используете ту же модель с другими параметрами, верно? Таким образом, пример команды может быть:

simple_hypersearch "python3 train_dnn.py --lr {lr} --batch_size {bs}" -p lr 0.001 0.0005 0.0001 -p bs 32 64 128 | simple_gpu_scheduler --gpus 0,1,2