From baf17892a7fcf62a0341ee5e1770cb6ad9adc083 Mon Sep 17 00:00:00 2001
From: KuNya <kunya.dn@gmail.com>
Date: Mon, 3 Sep 2018 19:37:44 +0800
Subject: [PATCH 1/2] add tensorboardX for loss visualization

---
 README.md               | 6 ++++++
 fastNLP/core/trainer.py | 4 ++++
 requirements.txt        | 1 +
 3 files changed, 11 insertions(+)

diff --git a/README.md b/README.md
index df9c89c9..c5e7a0c3 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,7 @@ fastNLP is a modular Natural Language Processing system based on PyTorch, for fa
 - numpy>=1.14.2
 - torch==0.4.0
 - torchvision>=0.1.8
+- tensorboardX
 
 
 ## Resources
@@ -45,6 +46,11 @@ conda install pytorch torchvision -c pytorch
 pip3 install torch torchvision
 ```
 
+### TensorboardX Installation
+
+```shell
+pip3 install tensorboardX
+```
 
 ## Project Structure
 
diff --git a/fastNLP/core/trainer.py b/fastNLP/core/trainer.py
index 7fc34da0..4b2e5248 100644
--- a/fastNLP/core/trainer.py
+++ b/fastNLP/core/trainer.py
@@ -5,6 +5,8 @@ import time
 from datetime import timedelta
 
 import torch
+import tensorboardX
+from tensorboardX import SummaryWriter
 
 from fastNLP.core.action import Action
 from fastNLP.core.action import RandomSampler, Batchifier
@@ -91,6 +93,7 @@ class BaseTrainer(object):
         self._loss_func = default_args["loss"].get()  # return a pytorch loss function or None
         self._optimizer = None
         self._optimizer_proto = default_args["optimizer"]
+        self._summary_writer = SummaryWriter(self.pickle_path + 'tensorboard_logs')
 
     def train(self, network, train_data, dev_data=None):
         """General Training Procedure
@@ -163,6 +166,7 @@ class BaseTrainer(object):
             loss = self.get_loss(prediction, batch_y)
             self.grad_backward(loss)
             self.update()
+            self._summary_writer.add_scalar("loss", loss.item(), global_step=step)
 
             if kwargs["n_print"] > 0 and step % kwargs["n_print"] == 0:
                 end = time.time()
diff --git a/requirements.txt b/requirements.txt
index d961dd92..954dd741 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 numpy>=1.14.2
 torch==0.4.0
 torchvision>=0.1.8
+tensorboardX

From 68b63fb071f7dd75be0fcdeb25251ec340785fad Mon Sep 17 00:00:00 2001
From: KuNya <kunya.dn@gmail.com>
Date: Tue, 4 Sep 2018 20:47:43 +0800
Subject: [PATCH 2/2] 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"]))