tensorflow指定GPU与动态分配GPU memory设置

下面我会详细讲解“TensorFlow指定GPU与动态分配GPU memory设置”的完整攻略。

指定GPU

在TensorFlow中，如果我们有多个GPU可用，可以通过以下两种方式指定哪个GPU要被使用：

1. 手动设置

可以手动设置环境变量CUDA_VISIBLE_DEVICES指定哪些GPU可被使用，这个环境变量的值是一个按逗号分隔的GPU编号列表，例如：

export CUDA_VISIBLE_DEVICES=0,1

上面的命令会将第一个和第二个GPU编号为0和1的GPU设为可见。

在Python中，可以使用以下代码手动设置GPU：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

这会将编号为0和1的GPU设为可见。

2. 使用函数指定

TensorFlow提供了两个函数可以用于指定使用哪个GPU：

• tf.ConfigProto
• tf.Session

我们可以通过在tf.ConfigProto对象中设置gpu_options属性指定使用哪个GPU，例如：

import tensorflow as tf

config = tf.ConfigProto()
config.gpu_options.visible_device_list = "0,1"
session = tf.Session(config=config)

这个代码会创建一个会话，GPU设备列表为编号为0和1的GPU。我们还可以在gpu_options属性中设置其他GPU属性，例如动态分配内存和允许内存增长。

动态分配GPU内存

默认情况下，TensorFlow将会分配GPU上所有空闲的内存，这种方式可以最大化地利用GPU资源，但是也可能导致GPU内存不足的错误。在这种情况下，我们可以使用以下两种方式限制TensorFlow使用的GPU内存量。

1. 手动限制

可以通过调用tf.GPUOptions类的per_process_gpu_memory_fraction参数指定GPU内存分配比例，例如：

import tensorflow as tf

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
session = tf.Session(config=config)

这个代码会创建一个会话，它将会限制TensorFlow使用50%的GPU内存。

2. 允许内存增长

我们也可以使用tf.GPUOptions类的allow_growth参数允许TensorFlow动态分配GPU内存，例如：

import tensorflow as tf

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

这个代码会创建一个会话，它将会允许TensorFlow动态分配GPU内存。

示例1：指定不同GPU

以下是一个指定不同GPU的完整示例：

import os
import tensorflow as tf

# 指定使用编号为1和2的GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

# 或者通过config指定
# config = tf.ConfigProto()
# config.gpu_options.visible_device_list = "1,2"
# session = tf.Session(config=config)

# 构建计算图
with tf.device('/gpu:1'):
    a = tf.placeholder(tf.float32)
    b = tf.placeholder(tf.float32)
    c = a + b

# 运行计算图
with tf.Session() as sess:
    print(sess.run(c, feed_dict={a: 3, b: 4}))

这个代码会使用编号为1和2的GPU构建计算图，并在运行时使用数量度来运行计算图。

示例2：限制GPU内存量

以下是一个限制GPU内存量的完整示例：

import tensorflow as tf

config = tf.ConfigProto()

# 指定使用GPU 0
config.gpu_options.visible_device_list = "0"

# 限制GPU内存使用量
config.gpu_options.per_process_gpu_memory_fraction = 0.4

# 或者允许内存增长
# config.gpu_options.allow_growth = True

# 构建计算图
a = tf.constant(1)
b = tf.constant(2)
c = a + b

# 运行计算图
with tf.Session(config=config) as sess:
    print(sess.run(c))

这个代码会在GPU 0上限制TensorFlow使用40%的GPU内存，并将计算图中的a + b操作运行在此GPU上。