鉴于有多个 GPU 可用，如何将专用 GPU 与 TF2 一起使用？

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.

鉴于有多个 GPU 可用，如何将专用 GPU 与 TF2 一起使用？

How to use dedicated GPU with TF2, given that multiple GPUs are available?

python

gpu

tensorflow2.0