《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,