Tensorflow 2.0知识点

文|Seraph

01 | Tensorflow2.0 GPU使用和分配：

1. 获取特定设备列表

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

2. 设置当前程序可见设备

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

3. 按需申请显存空间

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

4. 显存限制

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

5. 单GPU模拟多机GPU（测试下）

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