for gpu in tf.config.experimental.list_physical_devices(“GPU”):
tf.config.experimental.set_virtual_device_configuration(gpu,
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])
2. 仅在需要时获取GPU内存
for gpu in tf.config.experimental.list_physical_devices(“GPU”): tf.config.experimental.set_memory_growth(gpu, True)3. 将一个GPU分为两个虚拟设备
physical_gpus = tf.config.experimental.list_physical_devices(“GPU”)
tf.config.experimental.set_virtual_deivce_configuration(physical_gpus[0],
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])
参考文献
[1] A GéRon. Hands-on machine learning with Scikit-Learn and TensorFlow : concepts, tools, and techniques to build intelligent systems. 2017.
