tensorflow2 gpu使用

待补充图片和1.4

1. tensorflow2

默认情况下，tensorflow2会自动选择设备，如果有GPU，优先使用GPU，且自动占满多块GPU显存。

1.1 设置使用cpu

1. os.environ['CUDA_VISIBLE_DEVICES'] = "-1"
2. tf.config.experimental.set_visible_devices(devices=cpus[0], device_type='CPU')

1.2 设置使用GPU个数

首先可以先看看自己设备的GPU和CPU情况

# GPU
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
print(os.environ['CUDA_VISIBLE_DEVICES'])
# CPU
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')

• 如果有多块GPU，可以设置使用其中的几块，方法有三种：
  • 设置当前程序可见的设备范围：tf.config.experimental.set_visible_devices(devices=gpus[2:4], device_type='GPU')，设置设备2，3，4可见，当前程序只会使用自己可见的设备，不可见的设备不会被当前程序使用
  • 在终端使用环境变量限制GPU：export CUDA_VISIBLE_DEVICES=2,3
  • 在代码中指定GPU设备(类似于第一种方法)：os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"

1.3设置使用GPU显存大小

tensorflow提供了两种限定显存的方法：

1. 动态申请内存，需要多少，申请多少

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

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

# 限制使用2G显存
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)

1.4 单GPU模拟多GPU环境

当只有一块GPU时，假如显存为10G，但我们的程序可能只占用5G，那么剩下的5G不就浪费了，这时可以将这块10G的显存分成两块，这样相当于分布式使用2块GPU

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

按照上述，2块GPU也可以当做4块使用，具体方法待补充