It maintains the global step in the graph, making sure it's increased by one at every `hooked_sess.run`. This callback is used internally by the trainer, you don't need to worry about it.
| 103 | |
| 104 | |
| 105 | class MaintainStepCounter(Callback): |
| 106 | """ |
| 107 | It maintains the global step in the graph, making sure it's increased by one at every `hooked_sess.run`. |
| 108 | This callback is used internally by the trainer, you don't need to worry about it. |
| 109 | """ |
| 110 | |
| 111 | _chief_only = False |
| 112 | """ |
| 113 | In distributed training, we let each worker maintain its local global_step. |
| 114 | """ |
| 115 | |
| 116 | def _setup_graph(self): |
| 117 | # ensure it exists |
| 118 | gs_var = get_global_step_var() |
| 119 | with tf.name_scope(None): |
| 120 | self.gs_incr_op = tf.assign_add( |
| 121 | gs_var, 1, |
| 122 | name=GLOBAL_STEP_INCR_OP_NAME).op |
| 123 | self._fetches = tf.train.SessionRunArgs(self.gs_incr_op) |
| 124 | |
| 125 | def _before_train(self): |
| 126 | if self.global_step != 0: |
| 127 | logger.info("Start training with global_step={}".format(self.global_step)) |
| 128 | |
| 129 | def _before_run(self, _): |
| 130 | # always increase global_step when hooked_sess.run is called |
| 131 | return self._fetches |
| 132 | |
| 133 | def _after_run(self, _, __): |
| 134 | # Keep python-side global_step in agreement with TF-side |
| 135 | self.trainer.loop._global_step += 1 |
| 136 | |
| 137 | |
| 138 | class SessionRunTimeout(Callback): |