From 68b63fb071f7dd75be0fcdeb25251ec340785fad Mon Sep 17 00:00:00 2001 From: KuNya Date: Tue, 4 Sep 2018 20:47:43 +0800 Subject: [PATCH] add graph summary in _train_step --- fastNLP/core/trainer.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/fastNLP/core/trainer.py b/fastNLP/core/trainer.py index 4b2e5248..9e4bb215 100644 --- a/fastNLP/core/trainer.py +++ b/fastNLP/core/trainer.py @@ -94,6 +94,7 @@ class BaseTrainer(object): self._optimizer = None self._optimizer_proto = default_args["optimizer"] self._summary_writer = SummaryWriter(self.pickle_path + 'tensorboard_logs') + self._graph_summaried = False def train(self, network, train_data, dev_data=None): """General Training Procedure @@ -168,6 +169,10 @@ class BaseTrainer(object): self.update() self._summary_writer.add_scalar("loss", loss.item(), global_step=step) + if not self._graph_summaried: + self._summary_writer.add_graph(network, batch_x) + self._graph_summaried = True + if kwargs["n_print"] > 0 and step % kwargs["n_print"] == 0: end = time.time() diff = timedelta(seconds=round(end - kwargs["start"]))