| |
| from ...dist_utils import master_only |
| from ..hook import HOOKS |
| from .base import LoggerHook |
|
|
|
|
| @HOOKS.register_module() |
| class NeptuneLoggerHook(LoggerHook): |
| """Class to log metrics to NeptuneAI. |
| |
| It requires `neptune-client` to be installed. |
| |
| Args: |
| init_kwargs (dict): a dict contains the initialization keys as below: |
| - project (str): Name of a project in a form of |
| namespace/project_name. If None, the value of |
| NEPTUNE_PROJECT environment variable will be taken. |
| - api_token (str): User’s API token. |
| If None, the value of NEPTUNE_API_TOKEN environment |
| variable will be taken. Note: It is strongly recommended |
| to use NEPTUNE_API_TOKEN environment variable rather than |
| placing your API token in plain text in your source code. |
| - name (str, optional, default is 'Untitled'): Editable name of |
| the run. Name is displayed in the run's Details and in |
| Runs table as a column. |
| Check https://docs.neptune.ai/api-reference/neptune#init for |
| more init arguments. |
| interval (int): Logging interval (every k iterations). |
| ignore_last (bool): Ignore the log of last iterations in each epoch |
| if less than `interval`. |
| reset_flag (bool): Whether to clear the output buffer after logging |
| by_epoch (bool): Whether EpochBasedRunner is used. |
| |
| .. _NeptuneAI: |
| https://docs.neptune.ai/you-should-know/logging-metadata |
| """ |
|
|
| def __init__(self, |
| init_kwargs=None, |
| interval=10, |
| ignore_last=True, |
| reset_flag=True, |
| with_step=True, |
| by_epoch=True): |
|
|
| super(NeptuneLoggerHook, self).__init__(interval, ignore_last, |
| reset_flag, by_epoch) |
| self.import_neptune() |
| self.init_kwargs = init_kwargs |
| self.with_step = with_step |
|
|
| def import_neptune(self): |
| try: |
| import neptune.new as neptune |
| except ImportError: |
| raise ImportError( |
| 'Please run "pip install neptune-client" to install neptune') |
| self.neptune = neptune |
| self.run = None |
|
|
| @master_only |
| def before_run(self, runner): |
| if self.init_kwargs: |
| self.run = self.neptune.init(**self.init_kwargs) |
| else: |
| self.run = self.neptune.init() |
|
|
| @master_only |
| def log(self, runner): |
| tags = self.get_loggable_tags(runner) |
| if tags: |
| for tag_name, tag_value in tags.items(): |
| if self.with_step: |
| self.run[tag_name].log( |
| tag_value, step=self.get_iter(runner)) |
| else: |
| tags['global_step'] = self.get_iter(runner) |
| self.run[tag_name].log(tags) |
|
|
| @master_only |
| def after_run(self, runner): |
| self.run.stop() |
|
|
