Run an Op.
| 17 | |
| 18 | |
| 19 | class RunOp(Callback): |
| 20 | """ Run an Op. """ |
| 21 | |
| 22 | _chief_only = False |
| 23 | |
| 24 | def __init__(self, op, |
| 25 | run_before=True, run_as_trigger=True, |
| 26 | run_step=False, verbose=False): |
| 27 | """ |
| 28 | Args: |
| 29 | op (tf.Operation or function): an Op, or a function that returns the Op in the graph. |
| 30 | The function will be called after the main graph has been created (in the :meth:`setup_graph` callback). |
| 31 | run_before (bool): run the Op before training |
| 32 | run_as_trigger (bool): run the Op on every :meth:`trigger()` call. |
| 33 | run_step (bool): run the Op every step (along with training) |
| 34 | verbose (bool): print logs when the op is run. |
| 35 | |
| 36 | Example: |
| 37 | The `DQN Example |
| 38 | <https://github.com/tensorpack/tensorpack/blob/master/examples/DeepQNetwork/>`_ |
| 39 | uses this callback to update target network. |
| 40 | """ |
| 41 | if not callable(op): |
| 42 | self.setup_func = lambda: op # noqa |
| 43 | else: |
| 44 | self.setup_func = op |
| 45 | self.run_before = run_before |
| 46 | self.run_as_trigger = run_as_trigger |
| 47 | self.run_step = run_step |
| 48 | self.verbose = verbose |
| 49 | |
| 50 | def _setup_graph(self): |
| 51 | self._op = self.setup_func() |
| 52 | if self.run_step: |
| 53 | self._fetch = tf.train.SessionRunArgs(fetches=self._op) |
| 54 | |
| 55 | def _before_train(self): |
| 56 | if self.run_before: |
| 57 | self._print() |
| 58 | self._op.run() |
| 59 | |
| 60 | def _trigger(self): |
| 61 | if self.run_as_trigger: |
| 62 | self._print() |
| 63 | self._op.run() |
| 64 | |
| 65 | def _before_run(self, _): |
| 66 | if self.run_step: |
| 67 | self._print() |
| 68 | return self._fetch |
| 69 | |
| 70 | def _print(self): |
| 71 | if self.verbose: |
| 72 | logger.info("Running Op {} ...".format(self._op.name)) |
| 73 | |
| 74 | |
| 75 | class RunUpdateOps(RunOp): |
no outgoing calls
no test coverage detected