1.fix add_graph bug in SummaryWriter

2.replace tensorboardX by torch.utils.tensorboard
3.support for pytorch 1.2.0

jdit/trainer/super.py

@@ -12,11 +12,14 @@
 import torch
 import torchvision.transforms as transforms
 from torchvision.utils import make_grid
+
+
+from torch.utils.tensorboard import SummaryWriter
+
 import os
 import random
 import pandas as pd
 import numpy as np
-from tensorboardX import SummaryWriter
 
 from functools import wraps
 
@@ -662,7 +665,7 @@ def graph_lazy(self, model: Union[torch.nn.Module, torch.nn.DataParallel, Model]
 
         def hook(model, layer_input, layer_output):
             writer_for_model = SummaryWriter(log_dir=model_logdir)
-            input_for_test = tuple(i.detach().clone()[0] for i in layer_input)
+            input_for_test = tuple(i[0].detach().clone().unsqueeze(0) for i in layer_input)
             handel.remove()
             if isinstance(proto_model, torch.nn.DataParallel):
                 writer_for_model.add_graph(proto_model.module, input_for_test)

requirements.txt

@@ -10,6 +10,5 @@ pandas
 Pillow
 imageio
 scikit_learn
-tensorboardX
 imageio
 nvidia_ml_py3

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="jdit",  # pypi中的名称，pip或者easy_install安装时使用的名称，或生成egg文件的名称
-    version="0.0.17",
+    version="0.0.18",
     author="Guanglei Ding",
     author_email="dingguanglei.bupt@qq.com",
     maintainer='Guanglei Ding',
@@ -46,7 +46,7 @@
         "Programming Language :: Python :: 3 :: Only",
         'Programming Language :: Python :: Implementation',
         'Programming Language :: Python :: 3.6',
-        # 'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: Implementation :: CPython',
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
         'Topic :: Scientific/Engineering :: Visualization  '