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

Question

As title. I thought these lines would work to use only one 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')

But when I ran these following lines (I was following the tutorial from TF website):

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'])

The console still printed some lines showing that all 4 GPUs are allocated with some memory:

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

So what's the right way to restrict my GPU to only use the fourth(ie indexed [3] ) one?

---

Update : I also tried adding with ... , got same result:

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

These lines should be run first after those import s:

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

After these two commands, you also have to change this line from your OP:

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