Unverified Commit 9de7927a authored by Kai Chen's avatar Kai Chen Committed by GitHub
Browse files

Use torch.utils.tensorboard instead of tensorboardX (#108)

* use torch.utils.tensorboard instead of tensorboardX

* add backward compatibility

* add .idea to gitignore
parent 27ce4190
...@@ -103,6 +103,10 @@ venv.bak/ ...@@ -103,6 +103,10 @@ venv.bak/
# mypy # mypy
.mypy_cache/ .mypy_cache/
# editors and IDEs
.idea/
.vscode/
# custom # custom
mmcv/video/optflow_warp/flow_warp_module.cpp mmcv/video/optflow_warp/flow_warp_module.cpp
import os.path as osp import os.path as osp
import torch
from ...utils import master_only from ...utils import master_only
from .base import LoggerHook from .base import LoggerHook
...@@ -17,12 +19,20 @@ class TensorboardLoggerHook(LoggerHook): ...@@ -17,12 +19,20 @@ class TensorboardLoggerHook(LoggerHook):
@master_only @master_only
def before_run(self, runner): def before_run(self, runner):
if torch.__version__ >= '1.1':
try: try:
from tensorboardX import SummaryWriter from torch.utils.tensorboard import SummaryWriter
except ImportError: except ImportError:
raise ImportError('Please install tensorflow and tensorboardX ' raise ImportError(
'to use TensorboardLoggerHook.') 'Please run "pip install future tensorboard" to install '
'the dependencies to use torch.utils.tensorboard '
'(applicable to PyTorch 1.1 or higher)')
else: else:
try:
from tensorboardX import SummaryWriter
except ImportError:
raise ImportError('Please install tensorboardX to use '
'TensorboardLoggerHook.')
if self.log_dir is None: if self.log_dir is None:
self.log_dir = osp.join(runner.work_dir, 'tf_logs') self.log_dir = osp.join(runner.work_dir, 'tf_logs')
self.writer = SummaryWriter(self.log_dir) self.writer = SummaryWriter(self.log_dir)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment