| import os |
| import os.path as osp |
|
|
| from mmcv.runner import master_only, HOOKS, LoggerHook |
|
|
|
|
| @HOOKS.register_module() |
| class VisualDLLoggerHook(LoggerHook): |
|
|
| def __init__(self, |
| log_dir=None, |
| interval=10, |
| ignore_last=True, |
| reset_flag=True, |
| by_epoch=True): |
| super().__init__(interval, ignore_last, reset_flag, by_epoch) |
| afs_path = os.getenv("VDL_LOG_PATH") |
| if afs_path: |
| log_dir = afs_path.replace("afs://PUBLIC_KM_WD_Data:PUBLIC_km_wd_2020@wudang.afs.baidu.com:9902/user/PUBLIC_KM_WD_Data", 'afs') |
| self.log_dir = log_dir |
|
|
| @master_only |
| def before_run(self, runner): |
| try: |
| from visualdl import LogWriter |
| except ImportError: |
| raise ImportError('Please install visualdl to use ' |
| 'VisualDLLoggerHook.') |
|
|
| if self.log_dir is None: |
| self.log_dir = osp.join(runner.work_dir, 'VisualDl_logs') |
| else: |
| self.log_dir = osp.join(self.log_dir, runner.work_dir.split('/')[-1]) |
| print('VisualDL log dir: ', self.log_dir) |
| self.writer = LogWriter(self.log_dir) |
|
|
| @master_only |
| def log(self, runner): |
| tags = self.get_loggable_tags(runner, allow_text=True) |
| for tag, val in tags.items(): |
| if isinstance(val, str): |
| self.writer.add_text(tag, val, self.get_iter(runner)) |
| else: |
| self.writer.add_scalar(tag, val, self.get_iter(runner)) |
|
|
| @master_only |
| def after_run(self, runner): |
| self.writer.close() |
|
|
