allow period in mergeallsummaries

tensorpack/callbacks/summary.py

@@ -33,12 +33,18 @@ class MovingAverageSummary(Callback):
 
 
 class MergeAllSummaries_RunAlone(Callback):
-    def __init__(self, key):
+    def __init__(self, period, key):
+        self._period = period
         self._key = key
 
     def _setup_graph(self):
         self.summary_op = tf.summary.merge_all(self._key)
 
+    def _trigger_step(self):
+        if self._period:
+            if (self.local_step + 1) % self._period == 0:
+                self._trigger()
+
     def _trigger(self):
         if self.summary_op:
             summary = self.summary_op.eval()
@@ -46,7 +52,8 @@ class MergeAllSummaries_RunAlone(Callback):
 
 
 class MergeAllSummaries_RunWithOp(Callback):
-    def __init__(self, key):
+    def __init__(self, period, key):
+        self._period = period
         self._key = key
 
     def _setup_graph(self):
@@ -57,8 +64,15 @@ class MergeAllSummaries_RunWithOp(Callback):
             self._fetches = None
         self._total = self.trainer.config.steps_per_epoch
 
-    def _before_run(self, ctx):
+    def _need_run(self):
         if self.local_step == self._total - 1:
+            return True
+        if self._period > 0 and (self.local_step + 1) % self._period == 0:
+            return True
+        return False
+
+    def _before_run(self, ctx):
+        if self._need_run():
             return self._fetches
         return None
 
@@ -69,11 +83,13 @@ class MergeAllSummaries_RunWithOp(Callback):
         self.trainer.monitors.put_summary(summary)
 
 
-def MergeAllSummaries(run_alone=False, key=tf.GraphKeys.SUMMARIES):
+def MergeAllSummaries(period=0, run_alone=False, key=tf.GraphKeys.SUMMARIES):
     """
     Evaluate all summaries by `tf.summary.merge_all`, and write to logs.
 
     Args:
+        period (int): by default the callback summarizes once every epoch.
+            This option (if not set to 0) makes it additionally summarize every ``period`` steps.
         run_alone (bool): whether to evaluate the summaries alone.
             If True, summaries will be evaluated after each epoch alone.
             If False, summaries will be evaluated together with other
@@ -85,7 +101,8 @@ def MergeAllSummaries(run_alone=False, key=tf.GraphKeys.SUMMARIES):
     Returns:
         a Callback.
     """
+    period = int(period)
     if run_alone:
-        return MergeAllSummaries_RunAlone(key)
+        return MergeAllSummaries_RunAlone(period, key)
     else:
-        return MergeAllSummaries_RunWithOp(key)
+        return MergeAllSummaries_RunWithOp(period, key)