Tensorflow的GPU分配方法

1、查看服务器GPU内存使用情况命令行：

watch -n 1 nvidia-smi 

 图片显示的情况，页面每隔1秒会刷新页面。

2、Tensorflow使用多个GPU时的分配方法

方法一：在命令行中输入

CUDA_VISIBLE_DEVICES=0,1 python main.py

方法二：在代码中加入

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0，1"

即可指定程序只在显卡 0,1 上运行。

3、设置GPU占用量

import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用GPU90%的显存
sess = tf.Session(config=config)

4、配置GPU资源

TensorFlow 2.x 会自动检测环境中是否有GPU。如果有GPU，则会使用GPU。否则会使用CPU。

如果没有GPU，执行本部分也不会报错。

由于TensorFlow默认会申请全部GPU资源。为了达到GPU资源共享的目的，设置GPU资源按需申请

# 让GPU资源按需申请
import tensorflow as tf
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
 

5、获得当前主机上特定运算设备的列表

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

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'), PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')] [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

6、设置当前程序可见的设备范围

默认情况下 TensorFlow 会使用其所能够使用的所有 GPU

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

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

另一种方式是使用环境变量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的 GPU。

在终端输入

export CUDA_VISIBLE_DEVICES=0,1
# 或者在代码里加入
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"

都可以达到同样的效果。

.7、显存的使用

默认情况下，TensorFlow 将使用几乎所有可用的显存，以避免内存碎片化所带来的性能损失。

但是TensorFlow 提供两种显存使用策略，让我们能够更灵活地控制程序的显存使用方式：

    仅在需要时申请显存空间（程序初始运行时消耗很少的显存，随着程序的运行而动态申请显存）；

    限制消耗固定大小的显存（程序不会超出限定的显存大小，若超出的报错）。

设置仅在需要时申请显存空间。
 

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

下面的方式是设置Tensorflow固定消耗GPU:0的2GB显存。

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

8、单GPU模拟多GPU环境

上面的方式不仅可以设置显存的使用，还可以在只有单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)])

上面的代码就在GPU:0上建立了两个显存均为 2GB 的虚拟 GPU。