multi_gpu_model(num_gpus)
来生成多GPU
代码,并将对象保存在collection
中见之前博客构建TF代码
不要在单GPU代码中创建optimizer op
,因为是multi gpu
,所以参数更新的操作是所有的GPU
计算完梯度之后,才进行更新的。
multi_gpu_model
函数feed data
main 函数干了以下几件事:
1. 数据处理
2. 建立多GPU训练模型
3. 建立单/多GPU测试模型
4. 创建Saver
对象和FileWriter
对象
5. 创建session
6. run_epoch
还有一个版本,但是不work,不知为啥
# -*- coding:utf-8 -*-
import os.path
import re
import time
import numpy as np
import tensorflow as tf
import cifar10
batch_size=128
#train_dir='/tmp/cifar10_train'
#最大步数
max_steps=1000000
#GPU的数量
num_gpus=1
#log_device_placement=False
#计算损失的函数
def tower_loss(scope):
#使用cifar10.distorted_inputs()产生数据增强后的images和labels.
images, labels = cifar10.distorted_inputs()
#调用cifar10.inference生成卷积网络
#为每个GPU生成单独的网络,这些网络的结构完全一样,并且共享模型参数。
logits = cifar10.inference(images)
#调用cifar.loss计算损失函数(这里不直接返回loss,而是存储到collection中)
_ = cifar10.loss(logits, labels)
#使用tf.get_collection('losses', scope)获取当前这个GPU上的loss
#通过scope限定了范围
losses = tf.get_collection('losses', scope)
# 再使用tf.add_n将所有损失叠加到一起,得到total_loss,
total_loss = tf.add_n(losses, name='total_loss')
#返回total_loss作为函数结果
return total_loss
#它负责将不同GPU计算出的梯度进行合成
#合成的输入参数tower_grads是梯度的双层列表,
# 外层列表是不同GPU计算得到的梯度,内存列表是某个GPU内计算的不同Variable对应的梯度
#最内层元素为(grads,variable),即tower_grads的基本原始为二元组(梯度,变量),
#具体形式为:[[(grad0_gpu0,var0_gpu0),(grad1_gpu0,var1_gpu0)],[(grad0_gpu1,var0_gpu1),(grad1_gpu1,var1_gpu1)],....]]
#
#然后用循环便利这个双层列表
def average_gradients(tower_grads):
average_grads = []#首先创建平均梯度的列表average_grads,它负责将梯度在不同GPU间进行平均
#zip(*tower_grads)将这个双层列表转置,变成
#[[(grad0_gpu0,var0_gpu0),(grad0_gpu1,var0_gpu1)],[(grad1_gpu0,var1_gpu0),(grad1_gpu1,var1_gpu1)],....]]
#然后使用循环遍历其元素
for grad_and_vars in zip(*tower_grads):
# Note that each grad_and_vars looks like the following:
# ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
grads = []
for g, _ in grad_and_vars:
#每个循环获取的元素grad_and_vars,是同一个Variable的梯度在不同GPU计算的副本,需要计算其梯度的均值
#如果这个梯度是N维的向量,需要在每个维度上都进行平均。
# Add 0 dimension to the gradients to represent the tower.
#先使用tf.expand_dims给这些梯度添加一个冗余的维度0,然后把这些梯度都放到列表grad中。
expanded_g = tf.expand_dims(g, 0)
# Append on a 'tower' dimension which we will average over below.
#然后把这些梯度都放到列表grad中。
grads.append(expanded_g)
# Average over the 'tower' dimension.
#接着使用tf.concat将他们在维度0合并
grad = tf.concat(grads, 0)
#最后使用tf.reduce_mean针对维度0上求平均,即将其他维度全部平均。
grad = tf.reduce_mean(grad, 0)
#最后将平均后的梯度根Variable组合得到原有的二元组(梯度,变量)格式,并添加到列表average_grads
v = grad_and_vars[0][1]
grad_and_var = (grad, v)
average_grads.append(grad_and_var)
#当所有梯度都求完均值后,返回average_grads
return average_grads
def train():
#先设置默认的计算设备为CPU,用来计算一些简单的计算
with tf.Graph().as_default(), tf.device('/cpu:0'):
# Create a variable to count the number of train() calls. This equals the
# number of batches processed * FLAGS.num_gpus.
#然后使用global_step记录全局训练的步数
global_step = tf.get_variable(
'global_step', [],
initializer=tf.constant_initializer(0), trainable=False)
# Calculate the learning rate schedule.
#并计算一个epoch对应的batch数,
# 以及学习速率衰减需要的步数decay_steps.
num_batches_per_epoch = (cifar10.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN /
batch_size)
decay_steps = int(num_batches_per_epoch * cifar10.NUM_EPOCHS_PER_DECAY)
# Decay the learning rate exponentially based on the number of steps.
#使用tf.train.exponential_decay创建训练步数衰减的学习速率,
#第一个参数为初始学习速率
#第二个参数为全局训练的步数
#第三个参数为每次衰减需要的步数,
#第四个参数:为衰减率
#staircase设置为true,代表的是阶梯是的衰减,
lr = tf.train.exponential_decay(cifar10.INITIAL_LEARNING_RATE,
global_step,
decay_steps,
cifar10.LEARNING_RATE_DECAY_FACTOR,
staircase=True)
# Create an optimizer that performs gradient descent.
# 然后设置优化算法为GradientDescent,并传入随步数衰减的学习速率
opt = tf.train.GradientDescentOptimizer(lr)
# Calculate the gradients for each model tower.
#定义存储个GPU 计算结果的列表tower_grads
tower_grads = []
#创建一个循环,循环次数为GPU的数量
for i in range(num_gpus):
#在每一个循环内,使用tf.device限定使用第几个GPU,如gpu0,gpu1,
with tf.device('/gpu:%d' % i):
#然后使用tf.name_scope将命名空间定义为tower0,tower1的形式
with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:
#对每一个GPU,使用前面定义好的函数tower_loss获取其损失,
loss = tower_loss(scope)
# Reuse variables for the next tower.
#重用参数,让所有GPU共用一个模型以及完全相同的参数。
tf.get_variable_scope().reuse_variables()
# Calculate the gradients for the batch of data on this CIFAR tower.
#在使用opt.compute_gradients计算单个GPU的梯度
grads = opt.compute_gradients(loss)
# Keep track of the gradients across all towers.
#并将求得的梯度添加到梯度列表。
tower_grads.append(grads)
#使用前面写好的函数,average_gradients计算平均梯度
grads = average_gradients(tower_grads)
#并用opt.apply_gradients更新模型参数
apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
#这样就完成了多GPU的同步训练核参数更新
# Create a saver.
#创建模型的保存器saver.
saver = tf.train.Saver(tf.all_variables())
# Build an initialization operation to run below.
init = tf.global_variables_initializer()
#将Session的allow_soft_placement=True(有些操作在CPU上操作,不使用allow_soft_placement,会报错)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
#初始化全部参数
sess.run(init)
# Start the queue runners.准备号大量的数据增强后的训练样本,防止后妈的训练被阻塞在生成样本上。
#tf.train.start_queue_runners
tf.train.start_queue_runners(sess=sess)
#进行训练的循环,最大迭代次数max_steps。
for step in range(max_steps):
start_time = time.time()
#在每一步中执行一次更新梯度的操作apply_gradient_op(即一次训练操作)
_, loss_value = sess.run([apply_gradient_op, loss])
duration = time.time() - start_time#计算时间
# assert not np.isnan(loss_value), 'Model diverged with loss = NaN'
#每隔10步,展示一次当前batch的loss,以及每秒钟可训练的样本数核每个batch训练所需要话费的时间
if step % 10 == 0:
num_examples_per_step = batch_size * num_gpus
examples_per_sec = num_examples_per_step / duration
sec_per_batch = duration / num_gpus
format_str = ('step %d, loss = %.2f (%.1f examples/sec; %.3f '
'sec/batch)')
print (format_str % (step, loss_value,
examples_per_sec, sec_per_batch))
#每个1000步,使用saver保存整个模型文件
if step % 1000 == 0 or (step + 1) == max_steps:
# checkpoint_path = os.path.join(train_dir, 'model.ckpt')
saver.save(sess, './model/model.ckpt', global_step=step)
#使用cifar10.maybe_download_and_extract()下载完整的CIFAR-10数据,
cifar10.maybe_download_and_extract()
#调用train()函数开始训练
train()
联系客服