我可以在 tensorflow 2.x 中将物理 GPU 拆分为两个虚拟设备吗？

编程问答 2022-05-14

问题描述

我想测试使用多个 GPU 来训练网络。但是，我的机器上只安装了一个物理 GPU（RTX 2070）。我可以将此设备拆分为两个虚拟设备吗？
我目前的尝试是基于 tf.config.experimental.set_virtual_device_configuration 函数，但是，它似乎不起作用。完整示例：

import tensorflow as tf
print(tf.__version__) # 2.4.1
phisical_gpus = tf.config.experimental.list_physical_devices("GPU")
print(phisical_gpus) # [PhysicalDevice(name='/physical_device:GPU:0',device_type='GPU')]
tf.config.experimental.set_virtual_device_configuration(
    phisical_gpus[0],[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)
print(tf.config.experimental.list_physical_devices("GPU")) # [PhysicalDevice(name='/physical_device:GPU:0',device_type='GPU')]

解决方法

应使用 list_logical_devices 验证可见的 GPU 设备总数：

logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(logical_gpus)
#[LogicalDevice(name='/device:GPU:0',device_type='GPU'),# LogicalDevice(name='/device:GPU:1',device_type='GPU')]
gpugpunvidiatensorflowtensorflowtensorflowtensorflow2.x