我正在比较两个预训练模型，一个在 Tensorflow 中，一个在 Pytorch 中，在具有多个 GPU 的机器上。每个模型都适合一个 GPU。它们都加载在同一个 Python 脚本中。如何将一个 GPU 分配给 Tensorflow 模型，将另一个 GPU 分配给 Pytorch 模型？
设置 CUDA_VISIBLE_DEVICES=0,1 只会告诉两个模型这些 GPU 可用 - 我如何（我猜在 Python 中）确保 Tensorflow 使用 GPU 0 而 Pytorch 使用 GPU 1？
【问题讨论】：
标签： tensorflow cuda pytorch torch
您可以参考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
让 tensorflow 在 gpu:0 上加载模型，
with tf.device("gpu:0"):
load_model_function(model_path)
【讨论】：