class PaviLoggerHook(LoggerHook):
def __init__(self,
init_kwargs=None,
add_graph=False,
add_last_ckpt=False,
interval=10,
ignore_last=True,
reset_flag=True):
super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag)
self.init_kwargs = init_kwargs
self.add_graph = add_graph
self.add_last_ckpt = add_last_ckpt
@master_only
def before_run(self, runner):
try:
from pavi import SummaryWriter
except ImportError:
raise ImportError('Please run "pip install pavi" to install pavi.')
self.run_name = runner.work_dir.split('/')[-1]
if not self.init_kwargs:
self.init_kwargs = dict()
self.init_kwargs['task'] = self.run_name
self.init_kwargs['model'] = runner._model_name
self.writer = SummaryWriter(**self.init_kwargs)
if self.add_graph:
self.writer.add_graph(runner.model)
@master_only
def log(self, runner):
tags = {}
for tag, val in runner.log_buffer.output.items():
if tag not in ['time', 'data_time'] and is_scalar(val):
tags[tag] = val
tags['learning_rate'] = runner.current_lr()[0]
tags['momentum'] = runner.current_momentum()[0]
if tags:
self.writer.add_scalars(runner.mode, tags, runner.iter)
@master_only
def after_run(self, runner):
if self.add_last_ckpt:
ckpt_path = osp.join(runner.work_dir, 'latest.pth')
self.writer.add_snapshot_file(tag=self.run_name,
snapshot_file_path=ckpt_path,
iteration=runner.iter)
class PaviLoggerHook(LoggerHook):
def __init__(self,
init_kwargs=None,
add_graph=False,
add_last_ckpt=False,
interval=10,
ignore_last=True,
reset_flag=False,
by_epoch=True,
img_key='img_info'):
super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag,
by_epoch)
self.init_kwargs = init_kwargs
self.add_graph = add_graph
self.add_last_ckpt = add_last_ckpt
self.img_key = img_key
@master_only
def before_run(self, runner):
super(PaviLoggerHook, self).before_run(runner)
try:
from pavi import SummaryWriter
except ImportError:
raise ImportError('Please run "pip install pavi" to install pavi.')
self.run_name = runner.work_dir.split('/')[-1]
if not self.init_kwargs:
self.init_kwargs = dict()
self.init_kwargs['task'] = self.run_name
self.init_kwargs['model'] = runner._model_name
if runner.meta is not None:
if 'config_dict' in runner.meta:
config_dict = runner.meta['config_dict']
assert isinstance(
config_dict,
dict), ('meta["config_dict"] has to be of a dict, '
f'but got {type(config_dict)}')
elif 'config_file' in runner.meta:
config_file = runner.meta['config_file']
config_dict = dict(mmcv.Config.fromfile(config_file))
else:
config_dict = None
if config_dict is not None:
# 'max_.*iter' is parsed in pavi sdk as the maximum iterations
# to properly set up the progress bar.
config_dict = config_dict.copy()
config_dict.setdefault('max_iter', runner.max_iters)
# non-serializable values are first converted in
# mmcv.dump to json
config_dict = json.loads(
mmcv.dump(config_dict, file_format='json'))
session_text = yaml.dump(config_dict)
self.init_kwargs['session_text'] = session_text
self.writer = SummaryWriter(**self.init_kwargs)
def get_step(self, runner):
"""Get the total training step/epoch."""
if self.get_mode(runner) == 'val' and self.by_epoch:
return self.get_epoch(runner)
else:
return self.get_iter(runner)
@master_only
def log(self, runner):
tags = self.get_loggable_tags(runner, add_mode=False)
if tags:
self.writer.add_scalars(
self.get_mode(runner), tags, self.get_step(runner))
@master_only
def after_run(self, runner):
if self.add_last_ckpt:
ckpt_path = osp.join(runner.work_dir, 'latest.pth')
if osp.isfile(ckpt_path):
ckpt_path = osp.join(runner.work_dir, os.readlink(ckpt_path))
# runner.epoch += 1 has been done before `after_run`.
iteration = runner.epoch if self.by_epoch else runner.iter
return self.writer.add_snapshot_file(
tag=self.run_name,
snapshot_file_path=ckpt_path,
iteration=iteration)
@master_only
def before_epoch(self, runner):
if runner.epoch == 0 and self.add_graph:
if is_module_wrapper(runner.model):
_model = runner.model.module
else:
_model = runner.model
device = next(_model.parameters()).device
data = next(iter(runner.data_loader))
image = data[self.img_key][0:1].to(device)
with torch.no_grad():
self.writer.add_graph(_model, image)
class PaviLoggerHook(LoggerHook):
def __init__(self,
init_kwargs=None,
add_graph=False,
add_last_ckpt=False,
interval=10,
ignore_last=True,
reset_flag=True,
by_epoch=True):
super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag,
by_epoch)
self.init_kwargs = init_kwargs
self.add_graph = add_graph
self.add_last_ckpt = add_last_ckpt
@master_only
def before_run(self, runner):
try:
from pavi import SummaryWriter
except ImportError:
raise ImportError('Please run "pip install pavi" to install pavi.')
self.run_name = runner.work_dir.split('/')[-1]
if not self.init_kwargs:
self.init_kwargs = dict()
self.init_kwargs['task'] = self.run_name
self.init_kwargs['model'] = runner._model_name
if runner.meta is not None and 'config_dict' in runner.meta:
config_dict = runner.meta['config_dict'].copy()
# 'max_.*iter' is parsed in pavi sdk as the maximum iterations
# to properly set up the progress bar.
config_dict.setdefault('max_iter', runner.max_iters)
# non-serializable values are first converted in mmcv.dump to json
config_dict = json.loads(mmcv.dump(config_dict,
file_format='json'))
session_text = yaml.dump(config_dict)
self.init_kwargs['session_text'] = session_text
self.writer = SummaryWriter(**self.init_kwargs)
if self.add_graph:
self.writer.add_graph(runner.model)
@master_only
def log(self, runner):
tags = {}
for tag, val in runner.log_buffer.output.items():
if tag not in ['time', 'data_time'] and is_scalar(val):
tags[tag] = val
# add learning rate
lrs = runner.current_lr()
if isinstance(lrs, dict):
for name, value in lrs.items():
tags[f'learning_rate/{name}'] = value[0]
else:
tags['learning_rate'] = lrs[0]
# add momentum
momentums = runner.current_momentum()
if isinstance(momentums, dict):
for name, value in momentums.items():
tags[f'momentum/{name}'] = value[0]
else:
tags['momentum'] = momentums[0]
if tags:
if runner.mode == 'val':
mode = runner.mode
# runner.epoch += 1 has been done before val workflow
epoch = runner.epoch
else:
mode = 'train' if 'time' in runner.log_buffer.output else 'val'
epoch = runner.epoch + 1
if mode == 'val' and self.by_epoch:
self.writer.add_scalars(mode, tags, epoch)
else:
self.writer.add_scalars(mode, tags, runner.iter)
@master_only
def after_run(self, runner):
if self.add_last_ckpt:
ckpt_path = osp.join(runner.work_dir, 'latest.pth')
if osp.isfile(ckpt_path):
ckpt_path = osp.join(runner.work_dir, os.readlink(ckpt_path))
# runner.epoch += 1 has been done before `after_run`.
iteration = runner.epoch if self.by_epoch else runner.iter
return self.writer.add_snapshot_file(
tag=self.run_name,
snapshot_file_path=ckpt_path,
iteration=iteration)
class PaviLoggerHook(LoggerHook):
def __init__(self, init_kwargs=None, add_graph=False, add_last_ckpt=False, interval=10, ignore_last=True, reset_flag=True, by_epoch=True):
super(PaviLoggerHook, self).__init__(
interval, ignore_last, reset_flag, by_epoch)
self.init_kwargs = init_kwargs
self.add_graph = add_graph
self.add_last_ckpt = add_last_ckpt
@master_only
def before_run(self, runner):
try:
from pavi import SummaryWriter
except ImportError:
raise ImportError('Please run "pip install pavi" to install pavi.')
self.run_name = runner.work_dir.split('/')[(- 1)]
if (not self.init_kwargs):
self.init_kwargs = dict()
self.init_kwargs['task'] = self.run_name
self.init_kwargs['model'] = runner._model_name
if (runner.meta is not None):
if ('config_dict' in runner.meta):
config_dict = runner.meta['config_dict']
assert isinstance(config_dict, dict), ''.join(
['meta["config_dict"] has to be of a dict, but got ', '{}'.format(type(config_dict))])
elif ('config_file' in runner.meta):
config_file = runner.meta['config_file']
config_dict = dict(mmcv.Config.fromfile(config_file))
else:
config_dict = None
if (config_dict is not None):
config_dict = config_dict.copy()
config_dict.setdefault('max_iter', runner.max_iters)
config_dict = json.loads(
mmcv.dump(config_dict, file_format='json'))
session_text = yaml.dump(config_dict)
self.init_kwargs['session_text'] = session_text
self.writer = SummaryWriter(**self.init_kwargs)
if self.add_graph:
self.writer.add_graph(runner.model)
def get_step(self, runner):
'Get the total training step/epoch.'
if ((self.get_mode(runner) == 'val') and self.by_epoch):
return self.get_epoch(runner)
else:
return self.get_iter(runner)
@master_only
def log(self, runner):
tags = self.get_loggable_tags(runner, add_mode=False)
if tags:
self.writer.add_scalars(self.get_mode(
runner), tags, self.get_step(runner))
@master_only
def after_run(self, runner):
if self.add_last_ckpt:
ckpt_path = osp.join(runner.work_dir, 'latest.pth')
if osp.isfile(ckpt_path):
ckpt_path = osp.join(runner.work_dir, os.readlink(ckpt_path))
iteration = (runner.epoch if self.by_epoch else runner.iter)
return self.writer.add_snapshot_file(tag=self.run_name, snapshot_file_path=ckpt_path, iteration=iteration)
class PaviLoggerHook(LoggerHook):
def __init__(self,
init_kwargs=None,
add_graph=False,
add_last_ckpt=False,
interval=10,
ignore_last=True,
reset_flag=True,
by_epoch=True):
super(PaviLoggerHook, self).__init__(interval, ignore_last, reset_flag,
by_epoch)
self.init_kwargs = init_kwargs
self.add_graph = add_graph
self.add_last_ckpt = add_last_ckpt
@master_only
def before_run(self, runner):
try:
from pavi import SummaryWriter
except ImportError:
raise ImportError('Please run "pip install pavi" to install pavi.')
self.run_name = runner.work_dir.split('/')[-1]
if not self.init_kwargs:
self.init_kwargs = dict()
self.init_kwargs['task'] = self.run_name
self.init_kwargs['model'] = runner._model_name
if runner.meta is not None and 'config_dict' in runner.meta:
config_dict = runner.meta['config_dict'].copy()
# 'max_.*iter' is parsed in pavi sdk as the maximum iterations
# to properly set up the progress bar.
config_dict.setdefault('max_iter', runner.max_iters)
# non-serializable values are first converted in mmcv.dump to json
config_dict = json.loads(
mmcv.dump(config_dict, file_format='json'))
session_text = yaml.dump(config_dict)
self.init_kwargs['session_text'] = session_text
self.writer = SummaryWriter(**self.init_kwargs)
if self.add_graph:
self.writer.add_graph(runner.model)
@master_only
def log(self, runner):
tags = self.get_loggable_tags(runner)
if tags:
self.writer.add_scalars(
self.get_mode(runner), tags, self.get_step(runner))
@master_only
def after_run(self, runner):
if self.add_last_ckpt:
ckpt_path = osp.join(runner.work_dir, 'latest.pth')
if osp.isfile(ckpt_path):
ckpt_path = osp.join(runner.work_dir, os.readlink(ckpt_path))
# runner.epoch += 1 has been done before `after_run`.
iteration = runner.epoch if self.by_epoch else runner.iter
return self.writer.add_snapshot_file(
tag=self.run_name,
snapshot_file_path=ckpt_path,
iteration=iteration)