【TensorFlow】TensorFlow实战Google深度学习框架第2版参考代码（12-TensorFlow计算加速）

《TensorFlow实战Google深度学习框架第2版》教材中的样例代码，由于tensorflow版本、运行过程的修改、敲码过程中的失误、教材样例自带的错误等原因，可能会导致代码与教材不一致或者不能运行成功，仅供参考。

第12章 TensorFlow计算加速

本章中的内容涉及到TensorFlow集群，所以代码并没有进行运行。

• 12.3 在多GPU上训练深度学习模型解决MNIST问题，新的神经网络训练程序mnist_multi_gpu_train.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import os
import time
import tensorflow as tf
from datetime import datetime
from chapter05 import mnist_inference

# 定义训练神经网络时需要用到的参数
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 = 2

# 定义日志和模型输出的路径
MODEL_SAVE_PATH = "./logs_models/"
MODEL_NAME = "model.ckpt"

# 定义数据存储的路径
# 因为需要为不同的GPU提供不同的训练数据,所以通过placerholder的方式就需要手动准备多份数据
# 为了方便训练数据的获取过程,可以采用第7章中介绍的Dataset的方式从TFRecord中读取数据
# 于是在这里提供的数据文件路径为将MNIST训练数据转化为TFRecords格式之后的路径
DATA_PATH = "/path/to/output.tfrecords"


# 定义输入队列得到训练数据
def get_input():
    dataset = tf.data.TFRecordDataset([DATA_PATH])  # 这行代码有bug: dataset = tf.contrib.data.TFRecordDataset([DATA_PATH])

    # 定义数据解析格式
    def parser(record):
        features = tf.parse_single_example(
            record,
            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)

        return retyped_image, label

    # 定义输入队列
    dataset = dataset.map(parser)
    dataset = dataset.shuffle(buffer_size=10000)
    dataset = dataset.repeat(10)
    dataset = dataset.batch(BATCH_SIZE)
    iterator = dataset.make_one_shot_iterator()

    features, labels = iterator.get_next()
    return features, labels


# 定义损失函数
# 对于给定的训练数据、正则化损失计算规则和命名空间,计算在这个命名空间下的总损失
# 之所以需要给定命名空间是因为不同的GPU上计算得出的正则化损失都会加入名为loss的集合
# 如果不通过命名空间就会将不同GPU上的正则化损失都加进来
def get_loss(x, y_, regularizer, scope, reuse_variables=None):
    # 沿用5.5节中定义的函数来计算神经网络的前向传播结果
    with tf.variable_scope(tf.get_variable_scope(), reuse=reuse_variables):
        y = mnist_inference.inference(x, regularizer)
    # 计算交叉熵损失
    cross_entropy = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=y, labels=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(grads, 0)
        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'):
        # 定义基本的训练过程
        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,