| | |
| | from ...dist_utils import master_only |
| | from ..hook import HOOKS |
| | from .base import LoggerHook |
| |
|
| |
|
| | @HOOKS.register_module() |
| | class MlflowLoggerHook(LoggerHook): |
| |
|
| | def __init__(self, |
| | exp_name=None, |
| | tags=None, |
| | log_model=True, |
| | interval=10, |
| | ignore_last=True, |
| | reset_flag=False, |
| | by_epoch=True): |
| | """Class to log metrics and (optionally) a trained model to MLflow. |
| | |
| | It requires `MLflow`_ to be installed. |
| | |
| | Args: |
| | exp_name (str, optional): Name of the experiment to be used. |
| | Default None. |
| | If not None, set the active experiment. |
| | If experiment does not exist, an experiment with provided name |
| | will be created. |
| | tags (dict of str: str, optional): Tags for the current run. |
| | Default None. |
| | If not None, set tags for the current run. |
| | log_model (bool, optional): Whether to log an MLflow artifact. |
| | Default True. |
| | If True, log runner.model as an MLflow artifact |
| | for the current run. |
| | 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. |
| | |
| | .. _MLflow: |
| | https://www.mlflow.org/docs/latest/index.html |
| | """ |
| | super(MlflowLoggerHook, self).__init__(interval, ignore_last, |
| | reset_flag, by_epoch) |
| | self.import_mlflow() |
| | self.exp_name = exp_name |
| | self.tags = tags |
| | self.log_model = log_model |
| |
|
| | def import_mlflow(self): |
| | try: |
| | import mlflow |
| | import mlflow.pytorch as mlflow_pytorch |
| | except ImportError: |
| | raise ImportError( |
| | 'Please run "pip install mlflow" to install mlflow') |
| | self.mlflow = mlflow |
| | self.mlflow_pytorch = mlflow_pytorch |
| |
|
| | @master_only |
| | def before_run(self, runner): |
| | super(MlflowLoggerHook, self).before_run(runner) |
| | if self.exp_name is not None: |
| | self.mlflow.set_experiment(self.exp_name) |
| | if self.tags is not None: |
| | self.mlflow.set_tags(self.tags) |
| |
|
| | @master_only |
| | def log(self, runner): |
| | tags = self.get_loggable_tags(runner) |
| | if tags: |
| | self.mlflow.log_metrics(tags, step=self.get_iter(runner)) |
| |
|
| | @master_only |
| | def after_run(self, runner): |
| | if self.log_model: |
| | self.mlflow_pytorch.log_model(runner.model, 'models') |
| |
|
