實戰 Google 深度學習框架：TensorFlow 計算加速
阿新 • • 發佈：2018-12-14
# -*- coding: utf-8 -*-
from datetime import datetime
import os
import time
import tensorflow as tf
import mnist_inference
# 定義訓練神經網路時需要用到的配置。這些配置與5.5節中定義的配置類似。
BATCH_SIZE = 100
LEARNING_RATE_BASE = 0.001
LEARNING_RATE_DECAY = 0.99
REGULARAZTION_RATE = 0.0001
TRAINING_STEPS = 1000
MOVING_AVERAGE_DECAY = 0.99
N_GPU = 4
# 定義日誌和模型輸出的路徑。
MODEL_SAVE_PATH = "/path/to/logs_and_models/"
MODEL_NAME = "model.ckpt"
# 定義資料儲存的路徑。因為需要為不同的GPU提供不同的訓練資料，所以通過placerholder
# 的方式就需要手動準備多份資料。為了方便訓練資料的獲取過程，可以採用第7章中介紹的輸
# 入佇列的方式從TFRecord中讀取資料。於是在這裡提供的資料檔案路徑為將MNIST訓練資料
# 轉化為TFRecords格式之後的路徑。如何將MNIST資料轉化為TFRecord格式在第7章中有
# 詳細介紹，這裡不再贅述。
DATA_PATH = "/path/to/data.tfrecords"
# 定義輸入佇列得到訓練資料，具體細節可以參考第7章。
def get_input():
filename_queue = tf.train.string_input_producer([DATA_PATH])
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
# 定義資料解析格式。
features = tf.parse_single_example(
serialized_example,
features={
'image_raw': tf.FixedLenFeature([], tf.string),
'pixels': tf.FixedLenFeature([], tf.int64),
'label': tf.FixedLenFeature([], tf.int64),
})
# 解析圖片和標籤資訊。
decoded_image = tf.decode_raw(features['image_raw'], tf.uint8)
reshaped_image = tf.reshape(decoded_image, [784])
retyped_image = tf.cast(reshaped_image, tf.float32)
label = tf.cast(features['label'], tf.int32)
# 定義輸入佇列並返回。
min_after_dequeue = 10000
capacity = min_after_dequeue + 3 * BATCH_SIZE
return tf.train.shuffle_batch(
[retyped_image, label],
batch_size=BATCH_SIZE,
capacity=capacity,
min_after_dequeue=min_after_dequeue)
# 定義損失函式。對於給定的訓練資料、正則化損失計算規則和名稱空間，計算在這個名稱空間
# 下的總損失。之所以需要給定名稱空間是因為不同的GPU上計算得出的正則化損失都會加入名為
# loss的集合，如果不通過名稱空間就會將不同GPU上的正則化損失都加進來。
def get_loss(x, y_, regularizer, scope):
# 沿用5.5節中定義的函式來計算神經網路的前向傳播結果。
y = mnist_inference.inference(x, regularizer)
# 計算交叉熵損失。
cross_entropy = tf.reduce_mean(
tf.nn.sparse_softmax_cross_entropy_with_logits(y, y_))
# 計算當前GPU上計算得到的正則化損失。
regularization_loss = tf.add_n(tf.get_collection('losses', scope))
# 計算最終的總損失。
loss = cross_entropy + regularization_loss
return loss
# 計算每一個變數梯度的平均值。
def average_gradients(tower_grads):
average_grads = []
# 列舉所有的變數和變數在不同GPU上計算得出的梯度。
for grad_and_vars in zip(*tower_grads):
# 計算所有GPU上的梯度平均值。
grads = []
for g, _ in grad_and_vars:
expanded_g = tf.expand_dims(g, 0)
grads.append(expanded_g)
grad = tf.concat(0, grads)
grad = tf.reduce_mean(grad, 0)
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
# 將變數和它的平均梯度對應起來。
average_grads.append(grad_and_var)
# 返回所有變數的平均梯度，這將被用於變數更新。
return average_grads
# 主訓練過程。
def main(argv=None):
# 將簡單的運算放在CPU上，只有神經網路的訓練過程放在GPU上。
with tf.Graph().as_default(), tf.device('/cpu:0'):
# 獲取訓練batch。
x, y_ = get_input()
regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)
# 定義訓練輪數和指數衰減的學習率。
global_step = tf.get_variable(
'global_step', [], initializer=tf.constant_initializer(0),
trainable=False)
learning_rate = tf.train.exponential_decay(
LEARNING_RATE_BASE, global_step, 60000 / BATCH_SIZE,
LEARNING_ RATE_DECAY)
# 定義優化方法。
opt = tf.train.GradientDescentOptimizer(learning_rate)
tower_grads = []
# 將神經網路的優化過程跑在不同的GPU上。
for i in range(N_GPU):
# 將優化過程指定在一個GPU上。
with tf.device('/gpu:%d' % i):
with tf.name_scope('GPU_%d' % i) as scope:
cur_loss = get_loss(x, y_, regularizer, scope)
# 在第一次宣告變數之後，將控制變數重用的引數設定為True。這樣可以
# 讓不同的GPU更新同一組引數。注意tf.name_scope函式並不會影響
# tf.get_ variable的名稱空間。
tf.get_variable_scope().reuse_variables()
# 使用當前GPU計算所有變數的梯度。
grads = opt.compute_gradients(cur_loss)
tower_grads.append(grads)
# 計算變數的平均梯度，並輸出到TensorBoard日誌中。
grads = average_gradients(tower_grads)
for grad, var in grads:
if grad is not None:
tf.histogram_summary(
'gradients_on_average/%s' % var.op.name, grad)
# 使用平均梯度更新引數。
apply_gradient_op = opt.apply_gradients(
grads, global_step=global_ step)
for var in tf.trainable_variables():
tf.histogram_summary(var.op.name, var)
# 計算變數的滑動平均值。
variable_averages = tf.train.ExponentialMovingAverage(
MOVING_AVERAGE_DECAY, global_step)
variables_averages_op = variable_averages.apply(
tf.trainable_variables())
# 每一輪迭代需要更新變數的取值並更新變數的滑動平均值。
train_op = tf.group(apply_gradient_op, variables_averages_op)
saver = tf.train.Saver(tf.all_variables())
summary_op = tf.merge_all_summaries()
init = tf.initialize_all_variables()
# 訓練過程。
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True,
log_device_placement=True)) as sess:
# 初始化所有變數並啟動佇列。
init.run()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
summary_writer = tf.train.SummaryWriter(
MODEL_SAVE_PATH, sess.graph)
for step in range(TRAINING_STEPS):
# 執行神經網路訓練操作，並記錄訓練操作的執行時間。
start_time = time.time()
_, loss_value = sess.run([train_op, cur_loss])
duration = time.time() - start_time
# 每隔一段時間展示當前的訓練進度，並統計訓練速度。
if step != 0 and step % 10 == 0:
# 計算使用過的訓練資料個數。因為在每一次執行訓練操作時，每一個GPU
# 都會使用一個batch的訓練資料，所以總共用到的訓練資料個數為
# batch大小×GPU個數。
num_examples_per_step = BATCH_SIZE * N_GPU
# num_examples_per_step為本次迭代使用到的訓練資料個數，
# duration為運行當前訓練過程使用的時間，於是平均每秒可以處理的訓
# 練資料個數為num_examples_per_step / duration。
examples_per_sec = num_examples_per_step / duration
# duration為運行當前訓練過程使用的時間，因為在每一個訓練過程中，
# 每一個GPU都會使用一個batch的訓練資料，所以在單個batch上的訓
# 練所需要時間為duration / GPU個數。
sec_per_batch = duration / N_GPU
# 輸出訓練資訊。
format_str = ('step %d, loss = %.2f (%.1f examples/ '
' sec; %.3f sec/batch)')
print(format_str % (step, loss_value,
examples_per_sec, sec_per_batch))
# 通過TensorBoard視覺化訓練過程。
summary = sess.run(summary_op)
summary_writer.add_summary(summary, step)
# 每隔一段時間儲存當前的模型。
if step % 1000 == 0 or (step + 1) == TRAINING_STEPS:
checkpoint_path = os.path.join(
MODEL_SAVE_PATH, MODEL_ NAME)
saver.save(sess, checkpoint_path, global_step=step)
coord.request_stop()
coord.join(threads)
if __name__ == '__main__':
tf.app.run()
'''
在AWS的g2.8xlarge例項上執行上面這段程式可以得到類似下面的結果：
step 10, loss = 71.90 (15292.3 examples/sec; 0.007 sec/batch)
step 20, loss = 37.97 (18758.3 examples/sec; 0.005 sec/batch)
step 30, loss = 9.54 (16313.3 examples/sec; 0.006 sec/batch)
step 40, loss = 11.84 (14199.0 examples/sec; 0.007 sec/batch)
...
step 980, loss = 0.66 (15034.7 examples/sec; 0.007 sec/batch)
step 990, loss = 1.56 (16134.1 examples/sec; 0.006 sec/batch)
'''