tensorflow - GPU 加速

首先检测是否可用 GPU

import tensorflow as tf
print('GPU', tf.test.is_gpu_available())        # GPU True

tf.device

Tensorflow 通过 tf.device 指定每个操作运行的设备，可以指定本地的 CPU、GPU，还可以指定远程服务器；

Tensorflow 会给每个本地设备一个名称，如 /cpu:0，即使电脑有多块 CPU ，tf 不会做区分，统一叫 /cpu:0，而如果有多块 GPU，第 n 块 GPU 叫 /gpu:n，n 从 0 开始；

with tf.device('/cpu:0'):
    d1 = tf.Variable(1.)
    d2 = tf.Variable(2., name='d2')
with tf.device('/gpu:0'):
    d3 = tf.add(d1, d2)

with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(d3))

log_device_placement

tf 提供了 log_device_placement 来查看计算在哪个设备上运行；

d1 = tf.constant(1.)
d2 = tf.constant(2., name='d2')
d3 = tf.add(d1, d2)

### log_device_placement 记录了计算在哪个设备执行
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
    print(sess.run(d3))

# Device mapping:
# /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: Quadro P600, pci bus id: 0000:02:00.0, compute capability: 6.1
# Add: (Add): /job:localhost/replica:0/task:0/device:GPU:0
# Const: (Const): /job:localhost/replica:0/task:0/device:GPU:0
# d2: (Const): /job:localhost/replica:0/task:0/device:GPU:0
# 3.0

注意

1. 在配置好 GPU 的 tf 中，计算优先被分配到 GPU 上；

2. 如果有多块 GPU，tf 也会优先叫计算放到 /gpu:0 上，而其他 GPU 不会被安排任务，如果需要放到其他 GPU，可通过 tf.device 指定

allow_soft_placement

不是所有操作都能在 GPU 上运行；

如在 GPU 上用 tf.Variable 创建变量时，只支持实数型(float16、float32、double)，不支持整型；

# 在CPU上运行tf.Variable
a_cpu = tf.Variable(0, name="a_cpu")

with tf.device('/gpu:0'):
    # 将tf.Variable强制放在GPU上。
    # a_gpu = tf.Variable(0, name="a_gpu")
    ##### 上句报错如下
    # tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation a_gpu:
    # Could not satisfy explicit device specification '/device:GPU:0' because no supported kernel for GPU devices is available.
    a_gpu = tf.Variable(0., name="a_gpu")       ### 这样写就不报错

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

为了避免某些操作不能放在 GPU 上而报错，allow_soft_placement 可以将报错的操作自动放到 CPU 上；

a_cpu = tf.Variable(0, name="a_cpu")
with tf.device('/gpu:0'):
    a_gpu = tf.Variable(0, name="a_gpu")

# 通过allow_soft_placement参数自动将无法放在GPU上的操作放回CPU上
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True))
sess.run(tf.initialize_all_variables())


# Device mapping:
# /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: Quadro P600, pci bus id: 0000:02:00.0, compute capability: 6.1
# a_cpu: (VariableV2): /job:localhost/replica:0/task:0/device:CPU:0
# a_cpu/Assign: (Assign): /job:localhost/replica:0/task:0/device:CPU:0
# a_cpu/read: (Identity): /job:localhost/replica:0/task:0/device:CPU:0
# 
# a_gpu: (VariableV2): /job:localhost/replica:0/task:0/device:CPU:0
# a_gpu/Assign: (Assign): /job:localhost/replica:0/task:0/device:CPU:0
# a_gpu/read: (Identity): /job:localhost/replica:0/task:0/device:CPU:0        ##### 把 GPU 上的操作放到 CPU 上了
# 
# init/NoOp: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
# init/NoOp_1: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
# init: (NoOp): /job:localhost/replica:0/task:0/device:GPU:0
# a_cpu/initial_value: (Const): /job:localhost/replica:0/task:0/device:CPU:0
# a_gpu/initial_value: (Const): /job:localhost/replica:0/task:0/device:GPU:0

GPU 资源分配

虽然说 GPU 可以加速，但通常不会把所有操作都放在 GPU 上，大致有如下原则：

1. 把计算密集型的操作放到 GPU 上

GPU 是相对独立的资源，将计算转入和转出 GPU 都需要额外的时间，而且 GPU 需要将计算所需的数据从内存复制到 GPU 设备上，也需要额外的时间，

tensorflow 可自动完成这些操作，但为了提高运算效率，尽量将相关运算放到同一设备上；

2.通过环境变量分配 GPU 和显存；

################################### 通过环境变量分配 GPU 资源 ###################################
##### tensorflow 默认会占用所有 GPU (多块GPU) 和 所有显存，当然我们可以指定 GPU 或者 动态分配显存


################# 在运行时设置环境变量 #################
##### 只使用第二块GPU（GPU编号从0开始）
# 在demo_code.py中，机器上的第二块GPU的名称变成/gpu:0，在运行时所有/gpu:0的运算将被放在第二块GPU上
CUDA_VISIBLE_DEVICES=1 python demo_code.py

##### 只使用第一块和第二块GPU
CUDA_VISIBLE_DEVICES=0,1 python demo_code.py


################# 在程序中设置环境变量 #################
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'        ### 设定只使用 第3块 GPU


################# 动态分配显存 #################
##### TF 默认占用 GPU 的所有显存，我们可以手动分配显存，使得一块 GPU 可以同时运行多个任务
config = tf.ConfigProto()

# 让TensorFlow按需分配显存
config.gpu_options.allow_growth = True

# 或者直接按固定的比例分配
config.gpu_options.per_process_gpu_memory_fraction = 0.4        ### 占用所有可使用GPU的40%显存
session = tf.Session(config=config, ...)