| 85 | |
| 86 | |
| 87 | class MergeAllSummaries_RunWithOp(Callback): |
| 88 | def __init__(self, period, key): |
| 89 | self._period = period |
| 90 | self._key = key |
| 91 | |
| 92 | def _setup_graph(self): |
| 93 | size = len(tf.get_collection(self._key)) |
| 94 | logger.info("Summarizing collection '{}' of size {}.".format(self._key, size)) |
| 95 | self.summary_op = tf.summary.merge_all(self._key) |
| 96 | if self.summary_op is not None: |
| 97 | self._fetches = tf.train.SessionRunArgs(self.summary_op) |
| 98 | else: |
| 99 | self._fetches = None |
| 100 | |
| 101 | def _need_run(self): |
| 102 | if self.local_step == self.trainer.steps_per_epoch - 1: |
| 103 | return True |
| 104 | if self._period > 0 and (self.local_step + 1) % self._period == 0: |
| 105 | return True |
| 106 | return False |
| 107 | |
| 108 | def _before_run(self, ctx): |
| 109 | if self._need_run(): |
| 110 | return self._fetches |
| 111 | return None |
| 112 | |
| 113 | def _after_run(self, _, run_values): |
| 114 | summary = run_values.results |
| 115 | if summary is None: |
| 116 | return |
| 117 | self.trainer.monitors.put_summary(summary) |
| 118 | |
| 119 | |
| 120 | def MergeAllSummaries(period=0, run_alone=False, key=None): |