tensorflow使用多個gpu訓練

關於多gpu訓練,tf並無給太多的學習資料，比較官方的只有：tensorflow-models/tutorials/image/cifar10/cifar10_multi_gpu_train.py

但代碼比較簡單，只是針對cifar作了數據並行的多gpu訓練，利用到的layer、activation類型很少，針對更復雜網絡的狀況，並無給出指導。本身摸了很多坑以後，算是基本走通了，在此記錄下

 

1、思路

單GPU時，思路很簡單，前向、後向都在一個GPU上進行，模型參數更新時只涉及一個GPU。多GPU時，有模型並行和數據並行兩種狀況。模型並行指模型的不一樣部分在不一樣GPU上運行。數據並行指不一樣GPU上訓練數據不一樣，但模型是同一個（至關因而同一個模型的副本）。在此只考慮數據並行，這個在tf的實現思路以下：

模型參數保存在一個指定gpu/cpu上，模型參數的副本在不一樣gpu上，每次訓練，提供batch_size*gpu_num數據，並等量拆分紅多個batch，分別送入不一樣GPU。前向在不一樣gpu上進行，模型參數更新時，將多個GPU後向計算獲得的梯度數據進行平均，並在指定GPU/CPU上利用梯度數據更新模型參數。

假設有兩個GPU（gpu0,gpu1），模型參數實際存放在cpu0上，實際一次訓練過程以下圖所示：

 

2、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上。