| |
| 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') |
|
|
