CIFAR10 多 GPU 版本例程源码分析

1. # Copyright 2015 Google Inc. All Rights Reserved.  
2. #  
3. # Licensed under the Apache License, Version 2.0 (the "License");  
4. # you may not use this file except in compliance with the License.  
5. # You may obtain a copy of the License at  
6. #  
7. #     http://www.apache.org/licenses/LICENSE-2.0  
8. #  
9. # Unless required by applicable law or agreed to in writing, software  
10. # distributed under the License is distributed on an "AS IS" BASIS,  
11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
12. # See the License for the specific language governing permissions and  
13. # limitations under the License.  
14. # ==============================================================================  
15.   
16. """A binary to train CIFAR-10 using multiple GPU's with synchronous updates. 
17.  
18. Accuracy: 
19. cifar10_multi_gpu_train.py achieves ~86% accuracy after 100K steps (256 
20. epochs of data) as judged by cifar10_eval.py. 
21.  
22. Speed: With batch_size 128. 
23.  
24. System        | Step Time (sec/batch)  |     Accuracy 
25. -------------------------------------------------------------------- 
26. 1 Tesla K20m  | 0.35-0.60              | ~86% at 60K steps  (5 hours) 
27. 1 Tesla K40m  | 0.25-0.35              | ~86% at 100K steps (4 hours) 
28. 2 Tesla K20m  | 0.13-0.20              | ~84% at 30K steps  (2.5 hours) 
29. 3 Tesla K20m  | 0.13-0.18              | ~84% at 30K steps 
30. 4 Tesla K20m  | ~0.10                  | ~84% at 30K steps 
31.  
32. Usage: 
33. Please see the tutorial and website for how to download the CIFAR-10 
34. data set, compile the program and train the model. 
35.  
36. http://tensorflow.org/tutorials/deep_cnn/ 
37. """  
38. from __future__ import absolute_import  
39. from __future__ import division  
40. from __future__ import print_function  
41.   
42. from datetime import datetime  
43. import os.path  
44. import re  
45. import time  
46.   
47. import numpy as np  
48. from six.moves import xrange  # pylint: disable=redefined-builtin  
49. import tensorflow as tf  
50. from tensorflow.models.image.cifar10 import cifar10  
51.   
52. FLAGS = tf.app.flags.FLAGS  
53.   
54. tf.app.flags.DEFINE_string('train_dir', '/tmp/cifar10_train',  
55.                            """Directory where to write event logs """  
56.                            """and checkpoint.""")  
57. tf.app.flags.DEFINE_integer('max_steps', 1000000,  
58.                             """Number of batches to run.""")  
59. tf.app.flags.DEFINE_integer('num_gpus', 1,  
60.                             """How many GPUs to use.""")  
61. tf.app.flags.DEFINE_boolean('log_device_placement', False,  
62.                             """Whether to log device placement.""")  
63.   
64.   
65. def tower_loss(scope):  
66.   """Calculate the total loss on a single tower running the CIFAR model. 
67.  
68.   Args: 
69.     scope: unique prefix string identifying the CIFAR tower, e.g. 'tower_0' 
70.  
71.   Returns: 
72.      Tensor of shape [] containing the total loss for a batch of data 
73.   """  
74.   # Get images and labels for CIFAR-10.  
75.   images, labels = cifar10.distorted_inputs()  
76.   
77.   # Build inference Graph.  
78.   logits = cifar10.inference(images)  
79.   
80.   # Build the portion of the Graph calculating the losses. Note that we will  
81.   # assemble the total_loss using a custom function below.  
82.   _ = cifar10.loss(logits, labels)  
83.   
84.   # Assemble all of the losses for the current tower only.  
85.   losses = tf.get_collection('losses', scope)  
86.   
87.   # Calculate the total loss for the current tower.  
88.   total_loss = tf.add_n(losses, name='total_loss')  
89.   
90.   # Compute the moving average of all individual losses and the total loss.  
91.   loss_averages = tf.train.ExponentialMovingAverage(0.9, name='avg')  
92.   loss_averages_op = loss_averages.apply(losses + [total_loss])  
93.   
94.   # Attach a scalar summary to all individual losses and the total loss; do the  
95.   # same for the averaged version of the losses.  
96.   for l in losses + [total_loss]:  
97.     # Remove 'tower_[0-9]/' from the name in case this is a multi-GPU training  
98.     # session. This helps the clarity of presentation on tensorboard.  
99.     loss_name = re.sub('%s_[0-9]*/' % cifar10.TOWER_NAME, '', l.op.name)  
100.     # Name each loss as '(raw)' and name the moving average version of the loss  
101.     # as the original loss name.  
102.     tf.scalar_summary(loss_name +' (raw)', l)  
103.     tf.scalar_summary(loss_name, loss_averages.average(l))  
104.   
105.   with tf.control_dependencies([loss_averages_op]):  
106.     total_loss = tf.identity(total_loss)  
107.   return total_loss  
108.   
109.   
110. def average_gradients(tower_grads):  
111.   """Calculate the average gradient for each shared variable across all towers. 
112.  
113.   Note that this function provides a synchronization point across all towers. 
114.  
115.   Args: 
116.     tower_grads: List of lists of (gradient, variable) tuples. The outer list 
117.       is over individual gradients. The inner list is over the gradient 
118.       calculation for each tower. 
119.   Returns: 
120.      List of pairs of (gradient, variable) where the gradient has been averaged 
121.      across all towers. 
122.   """  
123.   average_grads = []  
124.   for grad_and_vars in zip(*tower_grads):  
125.     # Note that each grad_and_vars looks like the following:  
126.     #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))  
127.     grads = []  
128.     for g, _ in grad_and_vars:  
129.       # Add 0 dimension to the gradients to represent the tower.  
130.       expanded_g = tf.expand_dims(g, 0)  
131.   
132.       # Append on a 'tower' dimension which we will average over below.  
133.       grads.append(expanded_g)  
134.   
135.     # Average over the 'tower' dimension.  
136.     grad = tf.concat(0, grads)  
137.     grad = tf.reduce_mean(grad, 0)  
138.   
139.     # Keep in mind that the Variables are redundant because they are shared  
140.     # across towers. So .. we will just return the first tower's pointer to  
141.     # the Variable.  
142.     v = grad_and_vars[0][1]  
143.     grad_and_var = (grad, v)  
144.     average_grads.append(grad_and_var)  
145.   return average_grads  
146.   
147.   
148. def train():  
149.   """Train CIFAR-10 for a number of steps."""  
150.   with tf.Graph().as_default(), tf.device('/cpu:0'):  
151.     # Create a variable to count the number of train() calls. This equals the  
152.     # number of batches processed * FLAGS.num_gpus.  
153.     global_step = tf.get_variable(  
154.         'global_step', [],  
155.         initializer=tf.constant_initializer(0), trainable=False)  
156.   
157.     # Calculate the learning rate schedule.  
158.     num_batches_per_epoch = (cifar10.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN /  
159.                              FLAGS.batch_size)  
160.     decay_steps = int(num_batches_per_epoch * cifar10.NUM_EPOCHS_PER_DECAY)  
161.   
162.     # Decay the learning rate exponentially based on the number of steps.  
163.     lr = tf.train.exponential_decay(cifar10.INITIAL_LEARNING_RATE,  
164.                                     global_step,  
165.                                     decay_steps,  
166.                                     cifar10.LEARNING_RATE_DECAY_FACTOR,  
167.                                     staircase=True)  
168.   
169.     # Create an optimizer that performs gradient descent.  
170.     opt = tf.train.GradientDescentOptimizer(lr)  
171.   
172.     # Calculate the gradients for each model tower.  
173.     tower_grads = []  
174.     for i in xrange(FLAGS.num_gpus):  
175.       with tf.device('/gpu:%d' % i):  
176.         with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:  
177.           # Calculate the loss for one tower of the CIFAR model. This function  
178.           # constructs the entire CIFAR model but shares the variables across  
179.           # all towers.  
180.           loss = tower_loss(scope)  
181.   
182.           # Reuse variables for the next tower.  
183.           tf.get_variable_scope().reuse_variables()  
184.   
185.           # Retain the summaries from the final tower.  
186.           summaries = tf.get_collection(tf.GraphKeys.SUMMARIES, scope)  
187.   
188.           # Calculate the gradients for the batch of data on this CIFAR tower.  
189.           grads = opt.compute_gradients(loss)  
190.   
191.           # Keep track of the gradients across all towers.  
192.           tower_grads.append(grads)  
193.   
194.     # We must calculate the mean of each gradient. Note that this is the  
195.     # synchronization point across all towers.  
196.     grads = average_gradients(tower_grads)  
197.   
198.     # Add a summary to track the learning rate.  
199.     summaries.append(tf.scalar_summary('learning_rate', lr))  
200.   
201.     # Add histograms for gradients.  
202.     for grad, var in grads:  
203.       if grad is not None:  
204.         summaries.append(  
205.             tf.histogram_summary(var.op.name + '/gradients', grad))  
206.   
207.     # Apply the gradients to adjust the shared variables.  
208.     apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)  
209.   
210.     # Add histograms for trainable variables.  
211.     for var in tf.trainable_variables():  
212.       summaries.append(tf.histogram_summary(var.op.name, var))  
213.   
214.     # Track the moving averages of all trainable variables.  
215.     variable_averages = tf.train.ExponentialMovingAverage(  
216.         cifar10.MOVING_AVERAGE_DECAY, global_step)  
217.     variables_averages_op = variable_averages.apply(tf.trainable_variables())  
218.   
219.     # Group all updates to into a single train op.  
220.     train_op = tf.group(apply_gradient_op, variables_averages_op)  
221.   
222.     # Create a saver.  
223.     saver = tf.train.Saver(tf.all_variables())  
224.   
225.     # Build the summary operation from the last tower summaries.  
226.     summary_op = tf.merge_summary(summaries)  
227.   
228.     # Build an initialization operation to run below.  
229.     init = tf.initialize_all_variables()  
230.   
231.     # Start running operations on the Graph. allow_soft_placement must be set to  
232.     # True to build towers on GPU, as some of the ops do not have GPU  
233.     # implementations.  
234.     sess = tf.Session(config=tf.ConfigProto(  
235.         allow_soft_placement=True,  
236.         log_device_placement=FLAGS.log_device_placement))  
237.     sess.run(init)  
238.   
239.     # Start the queue runners.  
240.     tf.train.start_queue_runners(sess=sess)  
241.   
242.     summary_writer = tf.train.SummaryWriter(FLAGS.train_dir, sess.graph)  
243.   
244.     for step in xrange(FLAGS.max_steps):  
245.       start_time = time.time()  
246.       _, loss_value = sess.run([train_op, loss])  
247.       duration = time.time() - start_time  
248.   
249.       assert not np.isnan(loss_value), 'Model diverged with loss = NaN'  
250.   
251.       if step % 10 == 0:  
252.         num_examples_per_step = FLAGS.batch_size * FLAGS.num_gpus  
253.         examples_per_sec = num_examples_per_step / duration  
254.         sec_per_batch = duration / FLAGS.num_gpus  
255.   
256.         format_str = ('%s: step %d, loss = %.2f (%.1f examples/sec; %.3f '  
257.                       'sec/batch)')  
258.         print (format_str % (datetime.now(), step, loss_value,  
259.                              examples_per_sec, sec_per_batch))  
260.   
261.       if step % 100 == 0:  
262.         summary_str = sess.run(summary_op)  
263.         summary_writer.add_summary(summary_str, step)  
264.   
265.       # Save the model checkpoint periodically.  
266.       if step % 1000 == 0 or (step + 1) == FLAGS.max_steps:  
267.         checkpoint_path = os.path.join(FLAGS.train_dir, 'model.ckpt')  
268.         saver.save(sess, checkpoint_path, global_step=step)  
269.   
270.   
271. def main(argv=None):  # pylint: disable=unused-argument  
272.   cifar10.maybe_download_and_extract()  
273.   if tf.gfile.Exists(FLAGS.train_dir):  
274.     tf.gfile.DeleteRecursively(FLAGS.train_dir)  
275.   tf.gfile.MakeDirs(FLAGS.train_dir)  
276.   train()  
277.   
278.   
279. if __name__ == '__main__':  
280.   tf.app.run()