Как ограничить видимые устройства в TensorFlow? ⇐ Python

[Гость]

Я работаю на сервере с несколькими графическими процессорами и хочу выбрать один с достаточным количеством свободного места.
С помощью pytorch я обнаружил, что это можно сделать, просто установив для переменной среды ОС «CUDA_VISIBLE_DEVICES» значение нужные графические процессоры, но по какой-то причине тот же подход не работает с тензорным потоком.
Я нашел это решение:

Код: Выделить всё

config = tf.compat.v1.ConfigProto()
config.gpu_options.visible_device_list = "desired_GPU_no"
tf.compat.v1.keras.backend.set_session(tf.compat.v1.Session(config=config))

Unfortunately, it doesn't work for me, as running

Код: Выделить всё

tf.config.list_physical_devices('GPU')

shows all GPUs, not the ones I select.
Later on in my code, I build a tf.keras.layers.Embedding layer, and it fails because it automatically checks for space on GPU 0, which is occupied.
Are there any more reliable ways to choose the GPU I need, or am I doing something wrong with the current solution?
I should also specify that I am working on a jupyter notebook, maybe that is causing issues?


Источник: https://stackoverflow.com/questions/781 ... tensorflow