文|Seraph
01 | Tensorflow2.0 GPU使用和分配:
- 获取特定设备列表
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)
- 设置当前程序可见设备
tf.config.experimental.set_visible_devices(devices=gpus[2:4], device_type='GPU')
- 按需申请显存空间
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
- 显存限制
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)
- 单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)])