TensorFlow2.0——gpu设置

Tensorflow 2.0 GPU的使用与分配

或者参考TensorFlow中文文档 :讲述了各种情况

查看有什么设备

def set_soft_gpu(soft_gpu):
    import tensorflow as tf
    if soft_gpu:
        gpus = tf.config.experimental.list_physical_devices('GPU')
        if gpus:
            # Currently, memory growth needs to be the same across GPUs
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
            print(gpus, "Physical GPUs,", logical_gpus, "Logical GPUs")

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] Physical GPUs, 
[LogicalDevice(name='/device:GPU:0', device_type='GPU')] Logical GPUs

配置

虚拟gpu，一个gpu模拟多GPU
logical_gpus的作用:如果希望在不同的 GPU 上运行，则需要明确指定需要的 GPU：

#如果已经配置了gpu要重启：Virtual devices cannot be modified after being initialized
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  tf.config.experimental.set_virtual_device_configuration(
      gpus[0],
      [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
       tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])
  logical_gpus = tf.config.experimental.list_logical_devices('GPU')
  print(gpus, "Physical GPUs,", logical_gpus, "Logical GPUs")

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] Physical GPUs, 
[LogicalDevice(name='/device:GPU:0', device_type='GPU'),
 LogicalDevice(name='/device:GPU:1', device_type='GPU')] Logical GPUs

限制gpu使用

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

设置之后，当前程序只会使用自己可见的设备，不可见的设备不会被当前程序使用。

另一种方式是使用环境变量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的 GPU。
import os
os.environ[‘CUDA_VISIBLE_DEVICES’] = “2,3”

tf.config.list_logical_devices

英文不好、对机制不是很理解。只能附上英文解释

Logical devices may correspond to physical devices or remote devices in the cluster. Operations and tensors may be placed on these devices by using the name of the tf.config.LogicalDevice.

Calling tf.config.list_logical_devices triggers the runtime to configure any tf.config.PhysicalDevice visible to the runtime, thereby preventing further configuration. To avoid runtime initialization, call tf.config.list_physical_devices instead.

够用就行