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

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

0
JoSauderGH 11 Сен 2018 в 12:44

1 ответ

Лучший ответ

Вы можете обратиться к torch.device. https://pytorch.org/docs/stable/ tenor_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)

Поставить тензор / модель на gpu 0.

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

Для тензорного потока для загрузки модели на графический процессор: 0 do,

with tf.device("gpu:0"):
     load_model_function(model_path) 
1
Umang Gupta 13 Сен 2018 в 15:25