如何在张量流中自动选择空闲 GPU 进行模型训练？

Question

我正在使用 nvidia 预构建的 docker 容器NVIDIA Release 20.12-tf2来运行我的实验。我正在使用TensorFlow Version 2.3.1. 目前，我在其中一个 GPU 上运行我的模型，我还有 3 个空闲的 GPU，所以我打算在任何空闲的 GPU 上使用我的替代实验。这是输出nvidia-smi：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:6A:00.0 Off |                    0 |
| N/A   70C    P0    71W /  70W |  14586MiB / 15109MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla T4            Off  | 00000000:6B:00.0 Off |                    0 |
| N/A   39C    P0    27W /  70W |    212MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  Tesla T4            Off  | 00000000:6C:00.0 Off |                    0 |
| N/A   41C    P0    28W /  70W |    212MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  Tesla T4            Off  | 00000000:6D:00.0 Off |                    0 |
| N/A   41C    P0    28W /  70W |    212MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

更新：预建容器：

我正在使用 nvidia-prebuilt 容器，如下所示：

docker run -ti --rm --gpus all --shm-size=1024m -v /home/hamilton/data:/data nvcr.io/nvidia/tensorflow:20.12-tf2-py3

为了将空闲 GPU 用于我的其他实验，我尝试将它们添加到我的 python 脚本中：

尝试-1

import tensorflow as tf

devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(devices[0], True)

但这次尝试给了我以下错误：

raise ValueError("Memory growth cannot differ between GPU devices") ValueError: Memory growth cannot differ between GPU devices

我在谷歌上搜索了这个错误，但在 GitHub 上讨论的这些错误都不适合我。

尝试-2

我也试过这个：

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

但这种尝试也给了我这样的错误：

完成 GeneratorDataset 迭代器时出错：前提条件失败：Python 解释器状态未初始化。可以终止该过程。

人们在 github 上讨论了这个错误，但仍然无法摆脱我这边的错误。

最新尝试：

我还尝试使用 TensorFlow 进行并行训练，并将它们添加到我的 Python 脚本中：

device_type = "GPU"
devices = tf.config.experimental.list_physical_devices(device_type)
devices_names = [d.name.split("e:")[1] for d in devices]
strategy = tf.distribute.MirroredStrategy(devices=devices_names[:3])

with strategy.scope():
    opt = Adam(learning_rate=0.1)
    model.compile(loss='categorical_crossentropy', optimizer=opt, metrics=['accuracy'])

但这也给了我错误，程序停止了。

任何人都可以帮助我如何在 tensorflow 中为训练模型自动选择空闲 GPU？有谁知道任何可行的方法？我的上述尝试有什么问题？在其中一个 GPU 上运行程序时使用空闲 GPU 的任何可能想法？有什么想法吗？