GPU的分配及使用 TF2.0

默认情况下，TensorFlow 将使用几乎所有可用的显存，以避免内存碎片化所带来的性能损失。

但是TensorFlow 提供两种显存使用策略，让我们能够更灵活地控制程序的显存使用方式：

1.仅在需要时申请显存空间（程序初始运行时消耗很少的显存，随着程序的运行而动态申请显存）；

2.限制消耗固定大小的显存（程序不会超出限定的显存大小，若超出的报错）。

设置仅在需要时申请显存空间。

for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)

下面的方式是设置Tensorflow固定消耗GPU:0的4GB显存。

tf.config.experimental.set_virtual_device_configuration(
    gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]
)

获得当前主机上特定运算设备的列表

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)

设置当前程序可见的设备范围
默认情况下 TensorFlow 会使用其所能够使用的所有 GPU。

tf.config.experimental.set_visible_devices(devices=gpus[2:4], device_type='GPU')

设置之后，当前程序只会使用自己可见的设备，不可见的设备不会被当前程序使用。

另一种方式是使用环境变量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的 GPU。

在终端输入

export CUDA_VISIBLE_DEVICES=2,3

或者在代码里加入

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"

都可以达到同样的效果

参考：https://zhuanlan.zhihu.com/p/104173382
https://blog.csdn.net/zkbaba/article/details/108376987