關於多gpu訓練,tf並沒有給太多的學習資料，比較官方的只有：tensorflow-models/tutorials/image/cifar10/cifar10_multi_gpu_train.py
但代碼比較簡單，只是針對cifar做了數據並行的多gpu訓練，利用到的layer、activation類型不多，針對更復雜網絡的情況，並沒有給出指導。自己摸了不少坑之后，算是基本走通了，在此記錄下
一、思路
單GPU時，思路很簡單，前向、后向都在一個GPU上進行，模型參數更新時只涉及一個GPU。多GPU時，有模型並行和數據並行兩種情況。模型並行指模型的不同部分在不同GPU上運行。數據並行指不同GPU上訓練數據不同，但模型是同一個（相當於是同一個模型的副本）。在此只考慮數據並行，這個在tf的實現思路如下：
模型參數保存在一個指定gpu/cpu上，模型參數的副本在不同gpu上，每次訓練，提供batch_size*gpu_num數據，並等量拆分成多個batch，分別送入不同GPU。前向在不同gpu上進行，模型參數更新時，將多個GPU后向計算得到的梯度數據進行平均，並在指定GPU/CPU上利用梯度數據更新模型參數。
假設有兩個GPU（gpu0,gpu1），模型參數實際存放在cpu0上，實際一次訓練過程如下圖所示：
二、tf代碼實現
大部分需要修改的部分集中在構建計算圖上，假設在構建計算圖時，數據部分基於tensorflow1.4版本的dataset類，那么代碼要按照如下方式編寫：
1 next_img, next_label = iterator.get_next() 2 image_splits = tf.split(next_img, num_gpus) 3 label_splits = tf.split(next_label, num_gpus) 4 tower_grads = [] 5 tower_loss = [] 6 counter = 0 7 for d in self.gpu_id: 8 with tf.device('/gpu:%s' % d): 9 with tf.name_scope('%s_%s' % ('tower', d)): 10 cross_entropy = build_train_model(image_splits[counter], label_splits[counter], for_training=True) 11 counter += 1 12 with tf.variable_scope("loss"): 13 grads = opt.compute_gradients(cross_entropy) 14 tower_grads.append(grads) 15 tower_loss.append(cross_entropy) 16 tf.get_variable_scope().reuse_variables() 17 18 mean_loss = tf.stack(axis=0, values=tower_loss) 19 mean_loss = tf.reduce_mean(mean_loss, 0) 20 mean_grads = util.average_gradients(tower_grads) 21 update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) 22 with tf.control_dependencies(update_ops): 23 train_op = opt.apply_gradients(mean_grads, global_step=global_step)
第1行得到image和對應label
第2-3行對image和label根據使用的gpu數量做平均拆分（默認兩個gpu運算能力相同，如果gpu運算能力不同，可以自己設定拆分策略）
第 4-5行，保存來自不同GPU計算出的梯度、loss列表
第7-16行，開始在每個GPU上創建計算圖，最重要的是14-16三行，14，15把當前GPU計算出的梯度、loss值append到列表后，以便后續計算平均值。16行表示同名變量將會復用，這個是什么意思呢？假設現在gpu0上創建了兩個變量var0，var1，那么在gpu1上創建計算圖的時候，如果還有var0和var1，則默認復用之前gpu0上的創建的那兩個值。
第18-20行計算不同GPU獲取的grad、loss的平均值，其中第20行使用了cifar10_multi_gpu_train.py中的函數。
第23行利用梯度平均值更新參數。
注意：上述代碼中，所有變量(vars)都放在了第一個GPU上，運行時會發現第一個GPU占用的顯存比其他GPU多一些。如果想把變量放在CPU上，則需要在創建計算圖時，針對每層使用到的變量進行設備指定，很麻煩，所以建議把變量放在GPU上。