鉴于有多个 GPU 可用，如何在 TF2 中使用专用 GPU？

Question

如题。 我认为这些行只能使用一个 GPU：

_GPU = tf.config.list_physical_devices('GPU')[3]
tf.config.experimental.set_memory_growth(_GPU, True)
tf.config.set_visible_devices(_GPU, device_type='GPU')

但是当我运行以下这些行时（我正在关注 TF 网站上的教程）：

train_input_fn = make_input_fn(dftrain, y_train)
eval_input_fn = make_input_fn(dfeval, y_eval, num_epochs=1, shuffle=False)

linear_est = tf.estimator.LinearClassifier(feature_columns=feature_columns)
linear_est.train(train_input_fn)
result = linear_est.evaluate(eval_input_fn)

print(result['accuracy'])

控制台仍然打印了一些行，显示所有 4 个 GPU 都分配了一些内存：

INFO:tensorflow:Create CheckpointSaverHook.
INFO:tensorflow:Graph was finalized.
INFO:tensorflow:Running local_init_op.
2022-05-14 18:24:11.880130: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14874 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:b2:00.0, compute capability: 8.6
2022-05-14 18:24:11.880633: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 934 MB memory:  -> device: 1, name: NVIDIA GeForce RTX 3090, pci bus id: 0000:da:00.0, compute capability: 8.6
2022-05-14 18:24:11.881651: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:2 with 916 MB memory:  -> device: 2, name: NVIDIA RTX A5000, pci bus id: 0000:3d:00.0, compute capability: 8.6
2022-05-14 18:24:11.882376: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:3 with 4161 MB memory:  -> device: 3, name: NVIDIA RTX A5000, pci bus id: 0000:61:00.0, compute capability: 8.6

那么将我的 GPU 限制为仅使用第四个（即索引[3] ）的正确方法是什么？

---

更新：我也尝试添加with ... ，得到相同的结果：

with tf.device('/gpu:0'):

    linear_est = tf.estimator.LinearClassifier(feature_columns=feature_columns)
    linear_est.train(train_input_fn)
    result = linear_est.evaluate(eval_input_fn)

    print(result['accuracy'])

Answer 1

这些行应该在这些import之后首先运行：

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="2" # GPU according to nvidia-smi.

在这两个命令之后，您还必须从 OP 中更改此行：

tf.config.list_physical_devices('GPU')[0] # use [0] since now only 1 GPU.