TensorFlow中GPU的使用方法集合

1. 支持的设备

TensorFlow支持CPU和GPU这两种设备。它们均用字符串表示。例如：

• "/cpu:0": 机器的CPU
• "/device:GPU:0"：机器的第一个GPU（如果有一个）
• "/device:GPU:1"：机器的第二个GPU（以此类推）

如果TensorFlow指令中兼有CPU和GPU的实现，当该指令分配到设备时，GPU设备有优先权。例如，如果matmul同时存在CPU和GPU核函数，在同时有cpu:0和gpu:0设备的系统中gpu:0会被选来运行matmul。

 

2. 记录设备分配的方式

要找出指令和张量被分配到哪个设备，创建会话并将log_device_placement配置选项设为True。

# Creates a graph.
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

# Runs the op.
print(sess.run(c))

 

3. 手动分配设备

如果希望特定指令在选择的设备上运行，可以使用with tf.device 创建设备环境。在这个环境下，所有指令都将被分配在同一个设备上运行。

with tf.device('/cpu:0'):
  a = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[2,3], name='a')
  b = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[3,2], name='b')

c = tf.matmul(a, b)
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print(sess.run(c))

4. 允许增加GPU内存

        默认情况下，TensorFlow会映射进程可见的所有GPU的几乎所有GPU内存（还跟CUDA_VISIBLE_DEVICES有关）。通过减少内存碎片，可以更有效地使用设备上相对宝贵的GPU内存资源。在某些情况下，最理想的是进程只分配可用内存的一个子集，或者仅根据进程需要增加内存使用。TensorFlow在Session上提供两个Config选项来进行控制。

        第一个是allow_growth选项，它试图根据运行时的需要来分配GPU内存：它刚开始分配很少的内存，随着Session开始运行并需要更多的GPU内存，它会自动扩展TensorFlow进程所需的GPU内存区域。但是，它为了避免出现更加严重的内存碎片化情况并不会在运行时释放内存。要开启这选项，通过ConfigProto进行设置：

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

 

5. 在多GPU系统中使用单一GPU

如果系统中有多个GPU，则默认情况下将选择ID最小的GPU。如果希望在其他GPU上运行，则需要进行设置：

with tf.device('/device:GPU:2'):
  a = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[2,3], name='a')
  b = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[3,2], name='b')
  c = tf.matmul(a, b)

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print(sess.run( c ))

6. 使用多个GPU

如果要在多个GPU上运行TensorFlow，则可以采用多塔式方式构建模型，其中每个塔都会分配给不同GPU。例如：

c = []

for d in ['/device:GPU:2', '/device:GPU:3']:
  with tf.device(d):
    a = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[2,3])
    b = tf.constant([1.0,2.0,3.0,4.0,5.0,6.0], shape=[3,2])
    c.append(tf.matmul(a, b))

with tf.device('/cpu:0'):
  sum = tf.add_n( c)

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print(sess.run(sum))

我们需要用TensorFlow自带的函数或者类来建立计算图。用Python自有的运算符在没有重载的情况下将无法成功创建计算图。