From bbda6a8a6553202b8738523f74fc58658ff6ec4e Mon Sep 17 00:00:00 2001
From: chicm-ms <38930155+chicm-ms@users.noreply.github.com>
Date: Thu, 13 Aug 2020 09:59:10 +0800
Subject: [PATCH] Replace tensorboardX with torch.utils.tensorboard (#2786)

---
 azure-pipelines.yml                                          | 4 ----
 .../pynni/nni/compression/torch/pruning/amc/amc_pruner.py    | 5 ++---
 2 files changed, 2 insertions(+), 7 deletions(-)

diff --git a/azure-pipelines.yml b/azure-pipelines.yml
index 2889a8df67..78917d879f 100644
--- a/azure-pipelines.yml
+++ b/azure-pipelines.yml
@@ -28,7 +28,6 @@ jobs:
       set -e
       sudo apt-get install -y pandoc
       python3 -m pip install torch==1.5.0+cpu torchvision==0.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html --user
-      python3 -m pip install tensorboardX==1.9
       python3 -m pip install tensorflow==2.2.0 --user
       python3 -m pip install keras==2.4.2 --user
       python3 -m pip install gym onnx peewee thop --user
@@ -69,7 +68,6 @@ jobs:
   - script: |
       set -e
       python3 -m pip install torch==1.3.1+cpu torchvision==0.4.2+cpu -f https://download.pytorch.org/whl/torch_stable.html --user
-      python3 -m pip install tensorboardX==1.9
       python3 -m pip install tensorflow==1.15.2 --user
       python3 -m pip install keras==2.1.6 --user
       python3 -m pip install gym onnx peewee --user
@@ -119,7 +117,6 @@ jobs:
       set -e
       # pytorch Mac binary does not support CUDA, default is cpu version
       python3 -m pip install torchvision==0.6.0 torch==1.5.0 --user
-      python3 -m pip install tensorboardX==1.9
       python3 -m pip install tensorflow==1.15.2 --user
       brew install swig@3
       rm -f /usr/local/bin/swig
@@ -147,7 +144,6 @@ jobs:
       python -m pip install scikit-learn==0.23.2 --user
       python -m pip install keras==2.1.6 --user
       python -m pip install torch==1.5.0+cpu torchvision==0.6.0+cpu -f https://download.pytorch.org/whl/torch_stable.html --user
-      python -m pip install tensorboardX==1.9
       python -m pip install tensorflow==1.15.2 --user
     displayName: 'Install dependencies'
   - script: |
diff --git a/src/sdk/pynni/nni/compression/torch/pruning/amc/amc_pruner.py b/src/sdk/pynni/nni/compression/torch/pruning/amc/amc_pruner.py
index 5f3e1ce6be..2852fe6266 100644
--- a/src/sdk/pynni/nni/compression/torch/pruning/amc/amc_pruner.py
+++ b/src/sdk/pynni/nni/compression/torch/pruning/amc/amc_pruner.py
@@ -6,6 +6,7 @@
 from argparse import Namespace
 import numpy as np
 import torch
+from torch.utils.tensorboard import SummaryWriter
 
 from nni.compression.torch.compressor import Pruner
 from .channel_pruning_env import ChannelPruningEnv
@@ -148,8 +149,6 @@ def __init__(
             epsilon=50000,
             seed=None):
 
-        from tensorboardX import SummaryWriter
-
         self.job = job
         self.searched_model_path = searched_model_path
         self.export_path = export_path
@@ -189,7 +188,7 @@ def __init__(
 
         if self.job == 'train_export':
             print('=> Saving logs to {}'.format(self.output_dir))
-            self.tfwriter = SummaryWriter(logdir=self.output_dir)
+            self.tfwriter = SummaryWriter(log_dir=self.output_dir)
             self.text_writer = open(os.path.join(self.output_dir, 'log.txt'), 'w')
             print('=> Output path: {}...'.format(self.output_dir))