为Torch和Tensorflow分配两个独立的GPU

Question

我正在比较两个预先训练的模型，一个在Tensorflow中，另一个在Pytorch中，在具有多个GPU的机器上。每个型号都适合一个GPU。它们都加载在相同的Python脚本中。如何将一个GPU分配给Tensorflow模型，将另一个GPU分配给Pytorch模型？

设置CUDA_VISIBLE_DEVICES=0,1只告诉两个模型这些GPU是否可用 - 我怎么能（在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)

把张量/模型放在gpu 0上。

类似地，tensorflow具有tf.device。 https://www.tensorflow.org/api_docs/python/tf/device。它的用法在这里描述https://www.tensorflow.org/guide/using_gpu

for tensorflow to load model on gpu：0 do，

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