https://www.cnblogs.com/hrlnw/p/7779058.html

tensorflow使用多個gpu訓練

關於多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上。

 

分類: 技術, 研究 標籤: tensorflow, tf, gpu, 多gpu 好文要頂 關注我 收藏該文 handspeaker
關注 - 6
粉絲 - 234 +加關注 0 0 « 上一篇： tensorflow中moving average的用法
» 下一篇： centos系統安裝rar解壓工具unar

	</div>
	<div class="postDesc">posted @ <span id="post-date">2017-12-27 17:42</span> <a href="https://www.cnblogs.com/hrlnw/">handspeaker</a> 閱讀(<span id="post_view_count">5575</span>) 評論(<span id="post_comment_count">0</span>)  <a href="https://i.cnblogs.com/EditPosts.aspx?postid=7779058" rel="nofollow">編輯</a> <a href="#" onclick="AddToWz(7779058);return false;">收藏</a></div>
</div>