如何使用专用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]) 1 ?

---

更新:我也试着添加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

这些行应该在imports之后先运行:

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.