|
|
|
from ...dist_utils import master_only |
|
from ..hook import HOOKS |
|
from .base import LoggerHook |
|
|
|
|
|
@HOOKS.register_module() |
|
class WandbLoggerHook(LoggerHook): |
|
|
|
def __init__(self, |
|
init_kwargs=None, |
|
interval=10, |
|
ignore_last=True, |
|
reset_flag=False, |
|
commit=True, |
|
by_epoch=True, |
|
with_step=True): |
|
super(WandbLoggerHook, self).__init__(interval, ignore_last, |
|
reset_flag, by_epoch) |
|
self.import_wandb() |
|
self.init_kwargs = init_kwargs |
|
self.commit = commit |
|
self.with_step = with_step |
|
|
|
def import_wandb(self): |
|
try: |
|
import wandb |
|
except ImportError: |
|
raise ImportError( |
|
'Please run "pip install wandb" to install wandb') |
|
self.wandb = wandb |
|
|
|
@master_only |
|
def before_run(self, runner): |
|
super(WandbLoggerHook, self).before_run(runner) |
|
if self.wandb is None: |
|
self.import_wandb() |
|
if self.init_kwargs: |
|
self.wandb.init(**self.init_kwargs) |
|
else: |
|
self.wandb.init() |
|
|
|
@master_only |
|
def log(self, runner): |
|
tags = self.get_loggable_tags(runner) |
|
if tags: |
|
if self.with_step: |
|
self.wandb.log( |
|
tags, step=self.get_iter(runner), commit=self.commit) |
|
else: |
|
tags['global_step'] = self.get_iter(runner) |
|
self.wandb.log(tags, commit=self.commit) |
|
|
|
@master_only |
|
def after_run(self, runner): |
|
self.wandb.join() |
|
|
