Base 框架的 ModelFnHandler 对象添加 summary_op 参数用于自定义计算 summary 的操作

lienhua34 · lienhua34 · commit 3bc1ee9dc25c · 2017-04-25T16:55:28.000+08:00
&#25110;&#32773;&#20851;&#38381;&#33258;&#21160;&#35745;&#31639; summary &#26426;&#21046;&#12290;

&#40664;&#35748;&#24773;&#20917;&#19979;&#65292;Base &#26694;&#26550;&#20250;&#36890;&#36807; tf.summary.merge_all() &#20989;&#25968;&#26469;&#35745;&#31639;&#24182;&#25910;&#38598;&#27169;&#22411; Graph &#20013;
&#30340;&#25152;&#26377; summary &#20449;&#24687;&#12290;&#20294;&#26159;&#65292;&#24403;&#26576;&#20123; summary &#20449;&#24687;&#30340;&#35745;&#31639;&#38656;&#35201;&#23545; placeholder &#30340; feed &#26102;&#65292;
Base &#26694;&#26550;&#20250;&#21457;&#29983;&#24322;&#24120;&#12290;

&#20110;&#26159;&#36890;&#36807; ModelFnHandler &#23545;&#35937;&#30340; summary_op &#21442;&#25968;&#26469;&#35753;&#29992;&#25143;&#25552;&#20379;&#33258;&#21160;&#35745;&#31639; summary &#30340;&#25805;&#20316;&#65292;
&#25110;&#32773;&#23558;&#35813;&#21442;&#25968;&#35774;&#32622;&#20026; None &#26469;&#20851;&#38381;&#33258;&#21160;&#35745;&#31639;&#27169;&#22411; summary &#26426;&#21046;&#12290;
diff --git a/caicloud.tensorflow/caicloud/clever/examples/mnist/mnist_export.py b/caicloud.tensorflow/caicloud/clever/examples/mnist/mnist_export.py
@@ -33,12 +33,14 @@
 optimizer = None
 train_op = None
 global_step = None
+_summary_op = None
+_summary_writer = None
 
 def model_fn(sync, num_replicas):
     # &#36825;&#20123;&#21464;&#37327;&#22312;&#21518;&#32493;&#30340;&#35757;&#32451;&#25805;&#20316;&#20989;&#25968; train_fn() &#20013;&#20250;&#20351;&#29992;&#21040;&#65292;
     # &#25152;&#20197;&#36825;&#37324;&#20351;&#29992;&#20102; global &#21464;&#37327;&#12290;
     global input_images, loss, labels, optimizer, train_op, accuracy
-    global mnist, global_step
+    global mnist, global_step, _summary_op, _summary_writer
 
     # &#26500;&#24314;&#25512;&#29702;&#27169;&#22411;
     input_images = tf.placeholder(tf.float32, [None, 784], name='image')
@@ -54,6 +56,7 @@ def model_fn(sync, num_replicas):
     labels = tf.placeholder(tf.float32, [None, 10], name='labels')
     cross_entropy = tf.reduce_mean(
         tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels))
+    tf.summary.scalar("cross_entropy", cross_entropy)
     loss = tf.reduce_mean(cross_entropy, name='loss')
     tf.add_to_collection(tf.GraphKeys.LOSSES, loss)
         
@@ -68,6 +71,12 @@ def model_fn(sync, num_replicas):
             name="mnist_sync_replicas")
 
     train_op = optimizer.minimize(cross_entropy, global_step=global_step)
+
+    # &#33258;&#23450;&#20041;&#35745;&#31639;&#27169;&#22411; summary &#20449;&#24687;&#30340; Operation&#65292;
+    # &#24182;&#23450;&#20041;&#19968;&#20010; FileWriter &#29992;&#20110;&#20445;&#23384;&#27169;&#22411; summary &#20449;&#24687;&#12290;
+    # &#20854;&#20013;  dist_base.cfg.logdir &#26159; TaaS &#24179;&#21488;&#19978;&#35774;&#32622;&#30340;&#35757;&#32451;&#26085;&#24535;&#36335;&#24452;&#21442;&#25968;&#12290;
+    _summary_op = tf.summary.merge_all()
+    _summary_writer = tf.summary.FileWriter(dist_base.cfg.logdir)
         
     # Test trained model
     correct_prediction = tf.equal(tf.argmax(logits, 1),
@@ -89,12 +98,17 @@ def accuracy_evalute_fn(session):
     model_metric_ops = {
         "accuracy": accuracy_evalute_fn
     }
-    
+
+    # &#22240;&#20026;&#27169;&#22411;&#20013;&#38656;&#35201;&#35745;&#31639; tf.summary.scalar(cross_entropy)&#65292;&#32780;&#35813; summary &#30340;&#35745;&#31639;&#38656;&#35201;
+    # feed &#35774;&#32622; _input_images &#21644; _labels&#65292;&#25152;&#20197;&#36825;&#37324;&#23558; summary_op &#35774;&#32622;&#25104; None&#65292;&#23558;&#20851;&#38381;
+    # TaaS &#30340;&#33258;&#21160;&#35745;&#31639;&#21644;&#20445;&#23384;&#27169;&#22411; summary &#20449;&#24687;&#26426;&#21046;&#12290;&#22312; train_op &#20989;&#25968;&#20013;&#33258;&#24049;&#26469;&#35745;&#31639;&#24182;&#25910;&#38598;
+    # &#27169;&#22411; Graph &#30340; summary &#20449;&#24687;&#12290;
     return dist_base.ModelFnHandler(
         global_step=global_step,
         optimizer=optimizer,
         model_metric_ops = model_metric_ops,
-        model_export_spec=model_export_spec)
+        model_export_spec=model_export_spec,
+        summary_op=None)
     
 def gen_init_fn():
     """&#33719;&#21462;&#33258;&#23450;&#20041;&#21021;&#22987;&#21270;&#20989;&#25968;&#12290;
@@ -136,24 +150,34 @@ def init_from_checkpoint(scaffold, sess):
         print('Accuracy for restored model:')
         compute_accuracy(sess)
     return init_from_checkpoint
-    
+
+_last_summary_step = 0
 def train_fn(session, num_global_step):
     global local_step, input_images, labels, accuracy
     global mnist, train_op, loss, global_step
-    global local_step
+    global _summary_op, _summary_writer, _last_summary_step
     
     start_time = time.time()
     local_step += 1
     batch_xs, batch_ys = mnist.train.next_batch(100)
     feed_dict = {input_images: batch_xs,
                  labels: batch_ys}
-    _, loss_value, np_global_step = session.run(
-        [train_op, loss, global_step],
+    _, loss_value, np_global_step, summary_str = session.run(
+        [train_op, loss, global_step, _summary_op],
         feed_dict=feed_dict)
     duration = time.time() - start_time
     if local_step % 50 == 0:
         print('Step {0}: loss = {1:0.2f} ({2:0.3f} sec), global step: {3}.'.format(
             local_step, loss_value, duration, np_global_step))
+
+    # &#27599;&#38548;&#22266;&#23450;&#35757;&#32451;&#36718;&#25968;&#35745;&#31639;&#20445;&#23384;&#19968;&#27425;&#27169;&#22411; summary &#20449;&#24687;
+    # &#36890;&#36807; dist_base.cfg.save_summaies_steps &#33719;&#21462;&#22312; TaaS &#24179;&#21488;&#19978;&#35774;&#32622;&#30340;
+    # "&#33258;&#21160;&#20445;&#23384; summaries &#26085;&#24535;&#38388;&#38548;"&#21442;&#25968;&#20540;&#12290;
+    if (np_global_step - _last_summary_step >= dist_base.cfg.save_summaries_steps):
+        _summary_writer.add_summary(summary_str, np_global_step)
+        _summary_writer.flush()
+        _last_summary_step = np_global_step
+       
     if local_step % 1000 == 0:
         print("Accuracy for validation data: {0:0.3f}".format(
             session.run(
@@ -166,6 +190,9 @@ def train_fn(session, num_global_step):
         
         
 def after_train_hook(session):
+    global _summary_writer
+    _summary_writer.close()
+    
     print("Train done.")
     print("Accuracy for test data: {0:0.3f}".format(
         session.run(
diff --git a/caicloud.tensorflow/caicloud/clever/tensorflow/dist_base.py b/caicloud.tensorflow/caicloud/clever/tensorflow/dist_base.py
@@ -22,6 +22,8 @@
 from caicloud.clever.tensorflow import model_exporter
 import os
 
+_USE_DEFAULT = 0
+
 class RunConfig(object):
     def __init__(self):
         self.is_chief = True
@@ -41,7 +43,8 @@ def __init__(self,
                  global_step=None,
                  optimizer=None,
                  model_export_spec=None,
-                 model_metric_ops=None):
+                 model_metric_ops=None,
+                 summary_op=_USE_DEFAULT):
         """&#21019;&#24314;&#19968;&#20010; ModelFnHandler &#23545;&#35937;&#12290;
 
         &#20998;&#24067;&#24335; TensorFlow &#36816;&#34892;&#22120; DistTensorflowRunner &#30340;&#27169;&#22411;&#26500;&#24314;&#26041;&#27861; model_fn &#30340;&#36820;&#22238;&#20540;
@@ -73,6 +76,8 @@ def __init__(self,
         if (model_metric_ops is not None) and (not isinstance(model_metric_ops, dict)):
             raise ValueError('model_export_spec must be a None or dict.')
         self._model_metric_ops = model_metric_ops
+        
+        self._summary_op = summary_op
 
     @property
     def global_step(self):
@@ -90,6 +95,10 @@ def model_export_spec(self):
     def model_metric_ops(self):
         return self._model_metric_ops
 
+    @property
+    def summary_op(self):
+        return self._summary_op
+
 
 class DistTensorflowRunner(object):
     """&#20998;&#24067;&#24335; TensorFlow &#36816;&#34892;&#22120;&#12290;"""
@@ -170,6 +179,8 @@ def _call_model_fn(self):
         if model_fn_handler.model_export_spec is not None:
             self._model_exporter = model_exporter.ModelExporter(model_fn_handler.model_export_spec)
 
+        return model_fn_handler
+
     
     def run(self, train_fn):
         """&#25191;&#34892;&#20998;&#24067;&#24335; TensorFlow &#35757;&#32451;&#12290;
@@ -186,11 +197,13 @@ def run(self, train_fn):
 
         g = tf.Graph()
         with g.as_default():
-            self._call_model_fn()
+            model_fn_handler = self._call_model_fn()
             
 
             saver = tf.train.Saver()
-            summary_op = tf.summary.merge_all()
+            summary_op = model_fn_handler.summary_op
+            if summary_op == _USE_DEFAULT:
+                summary_op = tf.summary.merge_all()
             init_op = tf.global_variables_initializer()
 
             init_fn = None
