TensorFlow中GPU的使用小技巧

Table of Contents

1. tensorflow中表示设备的符号

2. 记录计算图结点的分配

3. 手动把结点分配到设备上

4. 允许增加GPU内存

5. 指定结点所在的GPU

6. 使用多个GPU

7. 允许程序弹性分配图结点

8. 分配一定比例的显存

---

---

1. tensorflow中表示设备的符号

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))

输出类似这样子

2018-12-04 21:22:50.908672: I T:\src\github\tensorflow\tensorflow\core\platform\cpu_feature_guard.cc:140] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2
2018-12-04 21:22:51.278386: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:1356] Found device 0 with properties: 
name: GeForce GTX 1060 6GB major: 6 minor: 1 memoryClockRate(GHz): 1.8095
pciBusID: 0000:01:00.0
totalMemory: 6.00GiB freeMemory: 4.96GiB
2018-12-04 21:22:51.278985: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:1435] Adding visible gpu devices: 0
2018-12-04 21:22:52.947815: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:923] Device interconnect StreamExecutor with strength 1 edge matrix:
2018-12-04 21:22:52.948151: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:929]      0 
2018-12-04 21:22:52.948411: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:942] 0:   N 
2018-12-04 21:22:52.948778: I T:\src\github\tensorflow\tensorflow\core\common_runtime\gpu\gpu_device.cc:1053] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 4732 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1060 6GB, pci bus id: 0000:01:00.0, compute capability: 6.1)

Device mapping: // 开始显示物理GPU到逻辑上device的映射
/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 1060 6GB, pci bus id: 0000:01:00.0, compute capability: 6.1
2018-12-04 21:22:53.151644: I T:\src\github\tensorflow\tensorflow\core\common_runtime\direct_session.cc:284] Device mapping:
/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce GTX 1060 6GB, pci bus id: 0000:01:00.0, compute capability: 6.1

// 开始显示计算结点到设备的映射

M2atMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0 // MatMul被映射到GPU:0
018-12-04 21:22:53.153550: I T:\src\github\tensorflow\tensorbflow\core\common_runtime\placer.cc:886] MatMul: (MatMul)/job:localhost/replica:0/task:0/device:GPU:0
2018-12-04: (Const): /job:localhost/replica:0/task:0/device:GPU:0
 21:22:53.153904: I T:\src\github\tensorflow\tensorflow\core\common_runtime\placer.cc:886] b: (Const)/job:localhost/replica:0/task:0/device:GPU:0 // b 被映射到GPU:0上
a: (Const): /job:localhost/replica:0/task:0/device:GPU:0 // a 被映射到GPU:0上
2018-12-04 21:22:53.154218: I T:\src\github\tensorflow\tensorflow\core\common_runtime\placer.cc:886] a: (Const)/job:localhost/replica:0/task:0/device:GPU:0 
// 计算结果
[[22. 28.]
 [49. 64.]]

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))

输出类似于这样子：

MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
2018-12-04 21:31:42.335148: I T:\src\github\tensorflow\tensorflow\core\common_runtime\placer.cc:886] MatMul: (MatMul)/job:localhost/replica:0/task:0/device:GPU:0 // 乘法操作分配到了GPU上
b2: (Const): /job:localhost/replica:0/task:0/device:CPU:0
018-12-04 21:31:42.335486: I T:\src\github\tensorflow\tensorflow\core\common_runtime\placer.cc:886] b: (Const)/job:localhost/replica:0/task:0/device:CPU:0
a: (Const): /job:localhost/replica:0/task:0/device:CPU:0
2018-12-04 21:31:42.335826: I T:\src\github\tensorflow\tensorflow\core\common_runtime\placer.cc:886] a: (Const)/job:localhost/replica:0/task:0/device:CPU:0  // a、b常量都在CPU上
[[22. 28.]
 [49. 64.]]

 

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，则默认情况下将选择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自有的运算符在没有重载的情况下将无法成功创建计算图。

7. 允许程序弹性分配图结点

如果程序中指定了计算图所在的设备，为了避免因找不到设备或者设备不支持结点操作而产生的错误，我们可以允许程序弹性分配图结点。

# Creates a graph.
with tf.device("/device:GPU:10"):
    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 allow_soft_placement set to True.
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

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

这样子在找不到GPU：10的情况下，或者GPU不支持我们需要的操作，程序会自动分配图结点到合适的设备上。

8. 分配一定比例的显存

gpu_options.per_process_gpu_memory_fraction 可以用来设置初始化时程序占用的显存比例。例如

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.1
config.gpu_options.allow_growth = True

可以在程序初始化时占用10%的显存。它与允许增加GPU内存一起使用，可以让程序只占有它所需的显存。为了避免显存碎片化问题，一旦显存被程序占用就不会被释放，直到程序运行结束。