tensorflow 單機多卡示例--數據並行

本文參考自官方的cifar10分類示例：
[url]https://www.tensorflow.org/tutorials/deep_cnn/[/url]

多機多卡（未驗證）：
[list]
[*][url]http://blog.csdn.net/cq361106306/article/details/52929468[/url]
[*][url]http://weibo.com/ttarticle/p/show?id=2309404005132982440427[/url]
[/list]

本文只保留了必要的代碼， 更適合於概念的理解。

在tensorflow中，變量是複用的，變量通過變量名唯一確定。
計算圖也會和設備綁定，如果一個圖計算時需要用到變量a，而變量a不在該設備上，則會自動生成相應的通信代碼，將變量a加載到該設備上。因而，變量的存放設備對於程序的正確性沒有影響，但會導致通信開銷有所差異。


測試結果： 對於全連接網絡，通信開銷佔比大，，，單卡最爲理想。。。
網絡大小：輸入2000*600， 中間層： 512, 128, 128, 1
運行時間：單位：秒
[img]http://dl2.iteye.com/upload/attachment/0122/4125/fbe13a1d-cfd4-3e7d-a430-9c8e29a74f09.png[/img]

# coding=utf-8
'''
Created on Jan 4, 2017
@author: colinliang

tensorflow 單機多卡程序示例， 
參考: tensorflow示例cifar10_multi_gpu_train.py
'''
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import numpy as np

def _allocate_variable(name, shape, initializer, dtype=tf.float32):
    # 分配變量，Tensorflow 會自動處理變量在不同設備間的通信問題，因而可以放在GPU上，也可以放在CPU上
    # 如果是單機單卡，都放在GPU上比較快 （無需顯式指定device, tf自動分配即可)
    # 如果是單機多卡，則放在CPU上略快；  可能是我這裏使用了SLI連接兩塊GPU，GPU間通信速度還算可以
    with tf.device('/cpu:0'): #強制放在主內存上
#     with tf.device(None): # 默認放在當前設備上
        var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype)
    print('%s: %s' % (var.op.name, var.device))
    return var

# 創建網絡 y=xw+b
def tower(input_tensor, target_tensor, scope, dims=[]):
    for i, d in enumerate(dims):
        with tf.variable_scope('affine%d' % i) as varscope:  # 僅僅用於生成變量的全名，與存放設備無關
            w = _allocate_variable('w', shape=[input_tensor.get_shape()[1], d], initializer=tf.truncated_normal_initializer(0, 1));
            b = _allocate_variable('b', shape=[], initializer=tf.zeros_initializer);
        input_tensor = tf.matmul(input_tensor, w) + b;
        input_tensor = tf.nn.relu(input_tensor)

    with tf.variable_scope('affine_last') as varscope:  # 僅僅用於生成變量的全名，與存放設備無關
#         w = _allocate_variable('w', shape=[input_tensor.get_shape()[1], 1], initializer=tf.truncated_normal_initializer(0, 1));
        w = _allocate_variable('w', shape=[input_tensor.get_shape()[1], 1], initializer=tf.constant_initializer(value=1));
        b = _allocate_variable('b', shape=[], initializer=tf.zeros_initializer);

    y = tf.matmul(input_tensor, w) + b;
    l = tf.reduce_mean(tf.square(y - target_tensor));
    tf.add_to_collection('losses', l)
    return y, l

# 合併所有tower上的梯度，取平均， 對於單機多卡程序，這段代碼是通用的
def average_tower_grads(tower_grads):
    print('towerGrads:')
    idx = 0
    for grads in tower_grads:  # grads 爲 一個list，其中元素爲 梯度-變量 組成的二元tuple
        print('grads---tower_%d' % idx)
        for g_var in grads:
            print(g_var)
            print('\t%s\n\t%s' % (g_var[0].op.name, g_var[1].op.name))
#             print('\t%s: %s'%(g_var[0].op.name,g_var[1].op.name))
        idx += 1

    if(len(tower_grads) == 1):
        return tower_grads[0]
    avgGrad_var_s = []
    for grad_var_s in zip(*tower_grads):
        grads = []
        v = None
        for g, v_ in grad_var_s:
            g = tf.expand_dims(g, 0)
            grads.append(g)
            v = v_
        all_g = tf.concat(0, grads)
        avg_g = tf.reduce_mean(all_g, 0, keep_dims=False)
        avgGrad_var_s.append((avg_g, v));
    return avgGrad_var_s

# 方案1 ，每組輸入分別用對應的placeholder作爲輸入;  未測試
def generate_towers_v1(NUM_GPU=2):  

    input_tensors = []
    target_tensors = []

    towerGrads = []
    lr = 1e-3
    opt = tf.train.AdamOptimizer(lr)

    for i in range(NUM_GPU):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:
                input_tensor = tf.placeholder(tf.float32, shape=[None, 1], name='input_%d' % i);
                input_tensors.append(input_tensor)
                target_tensor = tf.placeholder(tf.float32, shape=[None, 1], name='target_%d' % i);
                target_tensors.append(target_tensor)
                y, loss = tower(input_tensor=input_tensor, target_tensor=target_tensor, scope=scope)
                # Reuse variables for the next tower.
                tf.get_variable_scope().reuse_variables()
                grads = opt.compute_gradients(loss)
                towerGrads.append(grads)
    avgGrad_var_s = average_tower_grads(towerGrads)
    apply_gradient_op = opt.apply_gradients(avgGrad_var_s, global_step=None)
    loss = tf.Print(loss, data=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES))
    return input_tensors, target_tensors, y, loss, apply_gradient_op

# 方案2： 一組placeholder， 再根據tower數量分割成n組輸入，分別送人對應的tower
def generate_towers_v2(NUM_GPU=2, dim_in=1, dims=None, batch_size=None):   
    if(dims is None): dims = []

    input_tensor = tf.placeholder(tf.float32, shape=[batch_size, dim_in], name='input');
    target_tensor = tf.placeholder(tf.float32, shape=[batch_size, dim_in], name='target');
    input_tensors = tf.split(0, NUM_GPU, input_tensor)  # batch_size必須可以被dim_in整除
    target_tensors = tf.split(0, NUM_GPU, target_tensor)

    towerGrads = []
    lr = 1e-2
    opt = tf.train.AdamOptimizer(lr)  # 與GradientDescentOptimizer相比，會自動分配一些中間變量
    opt = tf.train.GradientDescentOptimizer(lr)
    for i in range(NUM_GPU):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:
                input_sub = input_tensors[i]
                print("device:%s" % input_sub.device)
                target_sub = target_tensors[i]
                y, loss = tower(input_tensor=input_sub, target_tensor=target_sub, scope=scope, dims=dims)
                # Reuse variables for the next tower.
                tf.get_variable_scope().reuse_variables()
                grads = opt.compute_gradients(loss)
                towerGrads.append(grads)
    avgGrad_var_s = average_tower_grads(towerGrads)
    loss = tf.Print(loss, data=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES))

    apply_gradient_op = opt.apply_gradients(avgGrad_var_s, global_step=None)

    print('ALL variables:')
    for v in tf.all_variables():
        print('\t%s' % v.op.name)

    return input_tensor, target_tensor, y, loss, apply_gradient_op

if __name__ == '__main__':
    sess = tf.Session()
    NUM_GPU = 1  # 由於只有兩塊GPU，如果設爲3，會報錯：Could not satisfy explicit device specification '/device:GPU:2'
    dim_in = 600; # 輸入變量x 的維度
    dims = [512, 128, 128] #隱層單元數，設置爲[]時表示 y=xw+b的線性變換，否則表示多層的全連接網絡
    batch_size = 2000; 

    input_tensor, target_tensor, y, loss, apply_gradient_op = generate_towers_v2(NUM_GPU=NUM_GPU, dim_in=dim_in, dims=dims)
    sess.run(tf.initialize_all_variables())

    inputs = np.random.rand(batch_size, dim_in)
    targets = inputs * 2 + 1;
    feed_dict = {input_tensor:inputs, target_tensor:targets}

    import time
    tstart = time.time()
    for i in range(10000):
#         _, l = sess.run([apply_gradient_op, loss], feed_dict=feed_dict)  #will print w, b
#         print(l)
        sess.run([apply_gradient_op], feed_dict=feed_dict)  # do not print w, b
    telapse = time.time() - tstart
    print(u'%d塊GPU用時: %.2fs' % (NUM_GPU, telapse))

示例輸出：
[quote]affine0/w: /device:CPU:0
affine0/b: /device:CPU:0
affine1/w: /device:CPU:0
affine1/b: /device:CPU:0
affine2/w: /device:CPU:0
affine2/b: /device:CPU:0
affine_last/w: /device:CPU:0
affine_last/b: /device:CPU:0
towerGrads:
grads---tower_0
(<tf.Tensor 'tower_0/gradients/tower_0/MatMul_grad/tuple/control_dependency_1:0' shape=(600, 512) dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c7144d0>)
tower_0/gradients/tower_0/MatMul_grad/tuple/control_dependency_1
affine0/w
(<tf.Tensor 'tower_0/gradients/tower_0/add_grad/tuple/control_dependency_1:0' shape=() dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c7140d0>)
tower_0/gradients/tower_0/add_grad/tuple/control_dependency_1
affine0/b
(<tf.Tensor 'tower_0/gradients/tower_0/MatMul_1_grad/tuple/control_dependency_1:0' shape=(512, 128) dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c7146d0>)
tower_0/gradients/tower_0/MatMul_1_grad/tuple/control_dependency_1
affine1/w
(<tf.Tensor 'tower_0/gradients/tower_0/add_1_grad/tuple/control_dependency_1:0' shape=() dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c6cb850>)
tower_0/gradients/tower_0/add_1_grad/tuple/control_dependency_1
affine1/b
(<tf.Tensor 'tower_0/gradients/tower_0/MatMul_2_grad/tuple/control_dependency_1:0' shape=(128, 128) dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c6cb750>)
tower_0/gradients/tower_0/MatMul_2_grad/tuple/control_dependency_1
affine2/w
(<tf.Tensor 'tower_0/gradients/tower_0/add_2_grad/tuple/control_dependency_1:0' shape=() dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c6f48d0>)
tower_0/gradients/tower_0/add_2_grad/tuple/control_dependency_1
affine2/b
(<tf.Tensor 'tower_0/gradients/tower_0/MatMul_3_grad/tuple/control_dependency_1:0' shape=(128, 1) dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c6f47d0>)
tower_0/gradients/tower_0/MatMul_3_grad/tuple/control_dependency_1
affine_last/w
(<tf.Tensor 'tower_0/gradients/tower_0/add_3_grad/tuple/control_dependency_1:0' shape=() dtype=float32>, <tensorflow.python.ops.variables.Variable object at 0x7f8b6c69f950>)
tower_0/gradients/tower_0/add_3_grad/tuple/control_dependency_1
affine_last/b
ALL variables:
affine0/w
affine0/b
affine1/w
affine1/b
affine2/w
affine2/b
affine_last/w
affine_last/b
[/quote]