为 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)

将 tensor/model 放在 gpu 0 上。

同样tensorflow有tf.device。 https://www.tensorflow.org/api_docs/python/tf/device. Its usage is described here https://www.tensorflow.org/guide/using_gpu

为了在 gpu:0 上加载模型，

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