Назначьте модели Torch и Tensorflow двумя отдельными графическими процессорами

Question

Я сравниваю две предварительно обученные модели, одна в Tensorflow и одна в Pytorch, на машине с несколькими графическими процессорами. Каждая модель умещается на одном графическом процессоре. Они оба загружены в одном и том же скрипте Python. Как я могу назначить один графический процессор модели Tensorflow и другой графический процессор модели Pytorch?

Установка CUDA_VISIBLE_DEVICES=0,1 только сообщает обеим моделям, что эти графические процессоры доступны - как я могу (в Python, я думаю), убедиться, что Tensorflow принимает GPU 0, а Pytorch принимает GPU 1?

Answer 1

Вы можете сослаться на torch.device. https://pytorch.org/docs/stable/tensor_attributes.html?highlight=device#torch.torch.device

В частности, сделать

device=torch.device("gpu:0")
tensor = tensor.to(device)

или для загрузки предварительно обученной модели

device=torch.device("gpu:0")
model = model.to(device)

чтобы положить тензор / модель на графический процессор 0.

Аналогично, у tenorflow есть tf.device. https://www.tensorflow.org/api_docs/python/tf/device. Его использование описано здесь https://www.tensorflow.org/guide/using_gpu

для тензорного потока для загрузки модели на GPU: 0 до,

with tf.device("gpu:0"):
     load_model_function(model_path)