def test_wrap_patch_with_module():
def new_log_param(key, value):
"""new mlflow.log_param"""
return (key, value)
before = get_func_attrs(mlflow.log_param)
wrap_patch(mlflow, mlflow.log_param.__name__, new_log_param)
after = get_func_attrs(mlflow.log_param)
assert after == before
assert mlflow.log_param("foo", "bar") == ("foo", "bar")
def autolog():
def __init__(self, *args, **kwargs):
original = gorilla.get_original_attribute(SparkSession, "__init__")
original(self, *args, **kwargs)
_listen_for_spark_activity(self._sc)
wrap_patch(SparkSession, "__init__", __init__)
active_session = _get_active_spark_session()
if active_session is not None:
# We know SparkContext exists here already, so get it
sc = SparkContext.getOrCreate()
_listen_for_spark_activity(sc)
def test_wrap_patch_with_module():
import sys
this_module = sys.modules[__name__]
def new_sample_function(a, b):
"""new mlflow.log_param"""
return a - b
before_attrs = get_func_attrs(mlflow.log_param)
assert sample_function_to_patch(10, 5) == 15
wrap_patch(this_module, sample_function_to_patch.__name__, new_sample_function)
after_attrs = get_func_attrs(mlflow.log_param)
assert after_attrs == before_attrs
assert sample_function_to_patch(10, 5) == 5
def test_wrap_patch_with_class():
class Math:
def add(self, a, b):
"""add"""
return a + b
def new_add(self, *args, **kwargs):
"""new add"""
orig = gorilla.get_original_attribute(self, "add")
return 2 * orig(*args, **kwargs)
before = get_func_attrs(Math.add)
wrap_patch(Math, Math.add.__name__, new_add)
after = get_func_attrs(Math.add)
assert after == before
assert Math().add(1, 2) == 6
def autolog():
"""
Enable automatic logging from Gluon to MLflow.
Logs loss and any other metrics specified in the fit
function, and optimizer data as parameters. Model checkpoints
are logged as artifacts to a 'models' directory.
"""
from mxnet.gluon.contrib.estimator import Estimator, EpochEnd, TrainBegin, TrainEnd
from mxnet.gluon.nn import HybridSequential
class __MLflowGluonCallback(EpochEnd, TrainEnd, TrainBegin):
def __init__(self):
self.current_epoch = 0
def epoch_end(self, estimator, *args, **kwargs):
logs = {}
for metric in estimator.train_metrics:
metric_name, metric_val = metric.get()
logs[metric_name] = metric_val
for metric in estimator.val_metrics:
metric_name, metric_val = metric.get()
logs[metric_name] = metric_val
try_mlflow_log(mlflow.log_metrics, logs, step=self.current_epoch)
self.current_epoch += 1
def train_begin(self, estimator, *args, **kwargs):
try_mlflow_log(mlflow.log_param, "num_layers", len(estimator.net))
if estimator.max_epoch is not None:
try_mlflow_log(mlflow.log_param, "epochs", estimator.max_epoch)
if estimator.max_batch is not None:
try_mlflow_log(mlflow.log_param, "batches",
estimator.max_batch)
try_mlflow_log(mlflow.log_param, "optimizer_name",
type(estimator.trainer.optimizer).__name__)
if hasattr(estimator.trainer.optimizer, "lr"):
try_mlflow_log(mlflow.log_param, "learning_rate",
estimator.trainer.optimizer.lr)
if hasattr(estimator.trainer.optimizer, "epsilon"):
try_mlflow_log(mlflow.log_param, "epsilon",
estimator.trainer.optimizer.epsilon)
def train_end(self, estimator, *args, **kwargs):
if isinstance(estimator.net, HybridSequential):
try_mlflow_log(log_model, estimator.net, artifact_path="model")
def fit(self, *args, **kwargs):
if not mlflow.active_run():
auto_end_run = True
else:
auto_end_run = False
original = gorilla.get_original_attribute(Estimator, "fit")
if len(args) >= 4:
l = list(args)
l[3] += [__MLflowGluonCallback()]
args = tuple(l)
elif "event_handlers" in kwargs:
kwargs["event_handlers"] += [__MLflowGluonCallback()]
else:
kwargs["event_handlers"] = [__MLflowGluonCallback()]
result = original(self, *args, **kwargs)
if auto_end_run:
mlflow.end_run()
return result
wrap_patch(Estimator, "fit", fit)
def autolog():
# pylint: disable=E0611
"""
Enables automatic logging from Keras to MLflow. Autologging captures the following information:
**Metrics** and **Parameters**
- Training loss; validation loss; user-specified metrics
- Metrics associated with the ``EarlyStopping`` callbacks: ``stopped_epoch``,
``restored_epoch``, ``restore_best_weight``, ``last_epoch``, etc
- ``fit()`` or ``fit_generator()`` parameters; optimizer name; learning rate; epsilon
- ``fit()`` or ``fit_generator()`` parameters associated with ``EarlyStopping``: ``min_delta``,
``patience``, ``baseline``, ``restore_best_weights``, etc
**Artifacts**
- Model summary on training start
- `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model) on training end
.. code-block:: python
:caption: Example
import mlflow
import mlflow.keras
# Build, compile, enable autologging, and train your model
keras_model = ...
keras_model.compile(optimizer="rmsprop", loss="mse", metrics=["accuracy"])
# autolog your metrics, parameters, and model
mlflow.keras.autolog()
results = keras_model.fit(
x_train, y_train, epochs=20, batch_size=128, validation_data=(x_val, y_val))
``EarlyStopping Integration with Keras AutoLogging``
MLflow will detect if an ``EarlyStopping`` callback is used in a ``fit()`` or
``fit_generator()`` call, and if the ``restore_best_weights`` parameter is set to be ``True``,
then MLflow will log the metrics associated with the restored model as a final, extra step.
The epoch of the restored model will also be logged as the metric ``restored_epoch``.
This allows for easy comparison between the actual metrics of the restored model and
the metrics of other models.
If ``restore_best_weights`` is set to be ``False``, then MLflow will not log an additional step.
Regardless of ``restore_best_weights``, MLflow will also log ``stopped_epoch``,
which indicates the epoch at which training stopped due to early stopping.
If training does not end due to early stopping, then ``stopped_epoch`` will be logged as ``0``.
MLflow will also log the parameters of the ``EarlyStopping`` callback,
excluding ``mode`` and ``verbose``.
"""
import keras
class __MLflowKerasCallback(keras.callbacks.Callback):
"""
Callback for auto-logging metrics and parameters.
Records available logs after each epoch.
Records model structural information as params when training begins
"""
def on_train_begin(self, logs=None): # pylint: disable=unused-argument
try_mlflow_log(mlflow.log_param, "num_layers",
len(self.model.layers))
try_mlflow_log(mlflow.log_param, "optimizer_name",
type(self.model.optimizer).__name__)
if hasattr(self.model.optimizer, "lr"):
lr = (self.model.optimizer.lr
if type(self.model.optimizer.lr) is float else
keras.backend.eval(self.model.optimizer.lr))
try_mlflow_log(mlflow.log_param, "learning_rate", lr)
if hasattr(self.model.optimizer, "epsilon"):
epsilon = (self.model.optimizer.epsilon
if type(self.model.optimizer.epsilon) is float else
keras.backend.eval(self.model.optimizer.epsilon))
try_mlflow_log(mlflow.log_param, "epsilon", epsilon)
sum_list = []
self.model.summary(print_fn=sum_list.append)
summary = "\n".join(sum_list)
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "model_summary.txt")
with open(summary_file, "w") as f:
f.write(summary)
try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def on_epoch_end(self, epoch, logs=None):
if not logs:
return
try_mlflow_log(mlflow.log_metrics, logs, step=epoch)
def on_train_end(self, logs=None):
try_mlflow_log(log_model, self.model, artifact_path="model")
# As of Keras 2.4.0, Keras Callback implementations must define the following
# methods indicating whether or not the callback overrides functions for
# batch training/testing/inference
def _implements_train_batch_hooks(self):
return False
def _implements_test_batch_hooks(self):
return False
def _implements_predict_batch_hooks(self):
return False
def _early_stop_check(callbacks):
if LooseVersion(keras.__version__) < LooseVersion("2.3.0"):
es_callback = keras.callbacks.EarlyStopping
else:
es_callback = keras.callbacks.callbacks.EarlyStopping
for callback in callbacks:
if isinstance(callback, es_callback):
return callback
return None
def _log_early_stop_callback_params(callback):
if callback:
try:
earlystopping_params = {
"monitor": callback.monitor,
"min_delta": callback.min_delta,
"patience": callback.patience,
"baseline": callback.baseline,
"restore_best_weights": callback.restore_best_weights,
}
try_mlflow_log(mlflow.log_params, earlystopping_params)
except Exception: # pylint: disable=W0703
return
def _get_early_stop_callback_attrs(callback):
try:
return callback.stopped_epoch, callback.restore_best_weights, callback.patience
except Exception: # pylint: disable=W0703
return None
def _log_early_stop_callback_metrics(callback, history):
if callback:
callback_attrs = _get_early_stop_callback_attrs(callback)
if callback_attrs is None:
return
stopped_epoch, restore_best_weights, patience = callback_attrs
try_mlflow_log(mlflow.log_metric, "stopped_epoch", stopped_epoch)
# Weights are restored only if early stopping occurs
if stopped_epoch != 0 and restore_best_weights:
restored_epoch = stopped_epoch - max(1, patience)
try_mlflow_log(mlflow.log_metric, "restored_epoch",
restored_epoch)
restored_metrics = {
key: history.history[key][restored_epoch]
for key in history.history.keys()
}
# Checking that a metric history exists
metric_key = next(iter(history.history), None)
if metric_key is not None:
last_epoch = len(history.history[metric_key])
try_mlflow_log(mlflow.log_metrics,
restored_metrics,
step=last_epoch)
def _run_and_log_function(self, original, args, kwargs, unlogged_params,
callback_arg_index):
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
log_fn_args_as_params(original, args, kwargs, unlogged_params)
early_stop_callback = None
# Checking if the 'callback' argument of the function is set
if len(args) > callback_arg_index:
tmp_list = list(args)
early_stop_callback = _early_stop_check(
tmp_list[callback_arg_index])
tmp_list[callback_arg_index] += [__MLflowKerasCallback()]
args = tuple(tmp_list)
elif "callbacks" in kwargs:
early_stop_callback = _early_stop_check(kwargs["callbacks"])
kwargs["callbacks"] += [__MLflowKerasCallback()]
else:
kwargs["callbacks"] = [__MLflowKerasCallback()]
_log_early_stop_callback_params(early_stop_callback)
history = original(self, *args, **kwargs)
_log_early_stop_callback_metrics(early_stop_callback, history)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return history
def fit(self, *args, **kwargs):
original = gorilla.get_original_attribute(keras.Model, "fit")
unlogged_params = [
"self", "x", "y", "callbacks", "validation_data", "verbose"
]
return _run_and_log_function(self, original, args, kwargs,
unlogged_params, 5)
def fit_generator(self, *args, **kwargs):
original = gorilla.get_original_attribute(keras.Model, "fit_generator")
unlogged_params = [
"self", "generator", "callbacks", "validation_data", "verbose"
]
return _run_and_log_function(self, original, args, kwargs,
unlogged_params, 4)
wrap_patch(keras.Model, "fit", fit)
wrap_patch(keras.Model, "fit_generator", fit_generator)
def autolog():
"""
Enables autologging for scikit-learn estimators.
**When is autologging performed?**
Autologging is performed when you call:
- ``estimator.fit()``
- ``estimator.fit_predict()``
- ``estimator.fit_transform()``
**Logged information**
**Parameters**
- Parameters obtained by ``estimator.get_params(deep=True)``. Note that ``get_params``
is called with ``deep=True``. This means when you fit a meta estimator that chains
a series of estimators, the parameters of these child estimators are also logged.
**Metrics**
- A training score obtained by ``estimator.score``. Note that the training score is
computed using parameters given to ``fit()``.
- Common metrics for classifier:
- `precision score`_
.. _precision score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.precision_score.html
- `recall score`_
.. _recall score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.recall_score.html
- `f1 score`_
.. _f1 score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.f1_score.html
- `accuracy score`_
.. _accuracy score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.accuracy_score.html
If the classifier has method ``predict_proba``, we additionally log:
- `log loss`_
.. _log loss:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.log_loss.html
- `roc auc score`_
.. _roc auc score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.roc_auc_score.html
- Common metrics for regressor:
- `mean squared error`_
.. _mean squared error:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html
- root mean squared error
- `mean absolute error`_
.. _mean absolute error:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_absolute_error.html
- `r2 score`_
.. _r2 score:
https://scikit-learn.org/stable/modules/generated/sklearn.metrics.r2_score.html
**Tags**
- An estimator class name (e.g. "LinearRegression").
- A fully qualified estimator class name
(e.g. "sklearn.linear_model._base.LinearRegression").
**Artifacts**
- A fitted estimator (logged by :py:func:`mlflow.sklearn.log_model()`).
**How does autologging work for meta estimators?**
When a meta estimator (e.g. `Pipeline`_, `GridSearchCV`_) calls ``fit()``, it internally calls
``fit()`` on its child estimators. Autologging does NOT perform logging on these constituent
``fit()`` calls.
**Parameter search**
In addition to recording the information discussed above, autologging for parameter
search meta estimators (`GridSearchCV`_ and `RandomizedSearchCV`_) records child runs
with metrics for each set of explored parameters, as well as artifacts and parameters
for the best model (if available).
**Supported estimators**
- All estimators obtained by `sklearn.utils.all_estimators`_ (including meta estimators).
- `Pipeline`_
- Parameter search estimators (`GridSearchCV`_ and `RandomizedSearchCV`_)
.. _sklearn.utils.all_estimators:
https://scikit-learn.org/stable/modules/generated/sklearn.utils.all_estimators.html
.. _Pipeline:
https://scikit-learn.org/stable/modules/generated/sklearn.pipeline.Pipeline.html
.. _GridSearchCV:
https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html
.. _RandomizedSearchCV:
https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.RandomizedSearchCV.html
**Example**
`See more examples <https://github.com/mlflow/mlflow/blob/master/examples/sklearn_autolog>`_
.. code-block:: python
from pprint import pprint
import numpy as np
from sklearn.linear_model import LinearRegression
import mlflow
def fetch_logged_data(run_id):
client = mlflow.tracking.MlflowClient()
data = client.get_run(run_id).data
tags = {k: v for k, v in data.tags.items() if not k.startswith("mlflow.")}
artifacts = [f.path for f in client.list_artifacts(run_id, "model")]
return data.params, data.metrics, tags, artifacts
# enable autologging
mlflow.sklearn.autolog()
# prepare training data
X = np.array([[1, 1], [1, 2], [2, 2], [2, 3]])
y = np.dot(X, np.array([1, 2])) + 3
# train a model
model = LinearRegression()
with mlflow.start_run() as run:
model.fit(X, y)
# fetch logged data
params, metrics, tags, artifacts = fetch_logged_data(run.info.run_id)
pprint(params)
# {'copy_X': 'True',
# 'fit_intercept': 'True',
# 'n_jobs': 'None',
# 'normalize': 'False'}
pprint(metrics)
# {'training_score': 1.0,
'training_mae': 2.220446049250313e-16,
'training_mse': 1.9721522630525295e-31,
'training_r2_score': 1.0,
'training_rmse': 4.440892098500626e-16}
pprint(tags)
# {'estimator_class': 'sklearn.linear_model._base.LinearRegression',
# 'estimator_name': 'LinearRegression'}
pprint(artifacts)
# ['model/MLmodel', 'model/conda.yaml', 'model/model.pkl']
"""
import pandas as pd
import sklearn
from mlflow.models import infer_signature
from mlflow.sklearn.utils import (
_MIN_SKLEARN_VERSION,
_is_supported_version,
_chunk_dict,
_get_args_for_score,
_log_specialized_estimator_content,
_get_Xy,
_all_estimators,
_truncate_dict,
_get_arg_names,
_get_estimator_info_tags,
_get_meta_estimators_for_autologging,
_is_parameter_search_estimator,
_log_parameter_search_results_as_artifact,
_create_child_runs_for_parameter_search,
)
from mlflow.tracking.context import registry as context_registry
from mlflow.utils.validation import (
MAX_PARAMS_TAGS_PER_BATCH,
MAX_PARAM_VAL_LENGTH,
MAX_ENTITY_KEY_LENGTH,
)
if not _is_supported_version():
warnings.warn(
"Autologging utilities may not work properly on scikit-learn < {} "
.format(_MIN_SKLEARN_VERSION) +
"(current version: {})".format(sklearn.__version__),
stacklevel=2,
)
def fit_mlflow(self, func_name, *args, **kwargs):
should_start_run = mlflow.active_run() is None
if should_start_run:
try_mlflow_log(mlflow.start_run)
_log_pretraining_metadata(self, *args, **kwargs)
original_fit = gorilla.get_original_attribute(self, func_name)
try:
fit_output = original_fit(*args, **kwargs)
except Exception as e:
if should_start_run:
try_mlflow_log(mlflow.end_run,
RunStatus.to_string(RunStatus.FAILED))
raise e
_log_posttraining_metadata(self, *args, **kwargs)
if should_start_run:
try_mlflow_log(mlflow.end_run)
return fit_output
def _log_pretraining_metadata(estimator, *args, **kwargs): # pylint: disable=unused-argument
"""
Records metadata (e.g., params and tags) for a scikit-learn estimator prior to training.
This is intended to be invoked within a patched scikit-learn training routine
(e.g., `fit()`, `fit_transform()`, ...) and assumes the existence of an active
MLflow run that can be referenced via the fluent Tracking API.
:param estimator: The scikit-learn estimator for which to log metadata.
:param args: The arguments passed to the scikit-learn training routine (e.g.,
`fit()`, `fit_transform()`, ...).
:param kwargs: The keyword arguments passed to the scikit-learn training routine.
"""
# Deep parameter logging includes parameters from children of a given
# estimator. For some meta estimators (e.g., pipelines), recording
# these parameters is desirable. For parameter search estimators,
# however, child estimators act as seeds for the parameter search
# process; accordingly, we avoid logging initial, untuned parameters
# for these seed estimators.
should_log_params_deeply = not _is_parameter_search_estimator(
estimator)
# Chunk model parameters to avoid hitting the log_batch API limit
for chunk in _chunk_dict(
estimator.get_params(deep=should_log_params_deeply),
chunk_size=MAX_PARAMS_TAGS_PER_BATCH,
):
truncated = _truncate_dict(chunk, MAX_ENTITY_KEY_LENGTH,
MAX_PARAM_VAL_LENGTH)
try_mlflow_log(mlflow.log_params, truncated)
try_mlflow_log(mlflow.set_tags, _get_estimator_info_tags(estimator))
def _log_posttraining_metadata(estimator, *args, **kwargs):
"""
Records metadata for a scikit-learn estimator after training has completed.
This is intended to be invoked within a patched scikit-learn training routine
(e.g., `fit()`, `fit_transform()`, ...) and assumes the existence of an active
MLflow run that can be referenced via the fluent Tracking API.
:param estimator: The scikit-learn estimator for which to log metadata.
:param args: The arguments passed to the scikit-learn training routine (e.g.,
`fit()`, `fit_transform()`, ...).
:param kwargs: The keyword arguments passed to the scikit-learn training routine.
"""
if hasattr(estimator, "score"):
try:
score_args = _get_args_for_score(estimator.score,
estimator.fit, args, kwargs)
training_score = estimator.score(*score_args)
except Exception as e: # pylint: disable=broad-except
msg = (
estimator.score.__qualname__ +
" failed. The 'training_score' metric will not be recorded. Scoring error: "
+ str(e))
_logger.warning(msg)
else:
try_mlflow_log(mlflow.log_metric, "training_score",
training_score)
# log common metrics and artifacts for estimators (classifier, regressor)
_log_specialized_estimator_content(estimator,
mlflow.active_run().info.run_id,
args, kwargs)
input_example = None
signature = None
if hasattr(estimator, "predict"):
try:
# Fetch an input example using the first several rows of the array-like
# training data supplied to the training routine (e.g., `fit()`)
fit_arg_names = _get_arg_names(estimator.fit)
X_var_name, y_var_name = fit_arg_names[:2]
input_example = _get_Xy(
args, kwargs, X_var_name,
y_var_name)[0][:INPUT_EXAMPLE_SAMPLE_ROWS]
model_output = estimator.predict(input_example)
signature = infer_signature(input_example, model_output)
except Exception as e: # pylint: disable=broad-except
input_example = None
msg = "Failed to infer an input example and model signature: " + str(
e)
_logger.warning(msg)
try_mlflow_log(
log_model,
estimator,
artifact_path="model",
signature=signature,
input_example=input_example,
)
if _is_parameter_search_estimator(estimator):
if hasattr(estimator, "best_estimator_"):
try_mlflow_log(
log_model,
estimator.best_estimator_,
artifact_path="best_estimator",
signature=signature,
input_example=input_example,
)
if hasattr(estimator, "best_score_"):
try_mlflow_log(mlflow.log_metric, "best_cv_score",
estimator.best_score_)
if hasattr(estimator, "best_params_"):
best_params = {
"best_{param_name}".format(param_name=param_name):
param_value
for param_name, param_value in
estimator.best_params_.items()
}
try_mlflow_log(mlflow.log_params, best_params)
if hasattr(estimator, "cv_results_"):
try:
# Fetch environment-specific tags (e.g., user and source) to ensure that lineage
# information is consistent with the parent run
environment_tags = context_registry.resolve_tags()
_create_child_runs_for_parameter_search(
cv_estimator=estimator,
parent_run=mlflow.active_run(),
child_tags=environment_tags,
)
except Exception as e: # pylint: disable=broad-except
msg = (
"Encountered exception during creation of child runs for parameter search."
" Child runs may be missing. Exception: {}".format(
str(e)))
_logger.warning(msg)
try:
cv_results_df = pd.DataFrame.from_dict(
estimator.cv_results_)
_log_parameter_search_results_as_artifact(
cv_results_df,
mlflow.active_run().info.run_id)
except Exception as e: # pylint: disable=broad-except
msg = (
"Failed to log parameter search results as an artifact."
" Exception: {}".format(str(e)))
_logger.warning(msg)
def patched_fit(self, func_name, *args, **kwargs):
"""
To be applied to a sklearn model class that defines a `fit` method and
inherits from `BaseEstimator` (thereby defining the `get_params()` method)
"""
with _SklearnTrainingSession(clazz=self.__class__,
allow_children=False) as t:
if t.should_log():
return fit_mlflow(self, func_name, *args, **kwargs)
else:
original_fit = gorilla.get_original_attribute(self, func_name)
return original_fit(*args, **kwargs)
def create_patch_func(func_name):
def f(self, *args, **kwargs):
return patched_fit(self, func_name, *args, **kwargs)
return f
_, estimators_to_patch = zip(*_all_estimators())
# Ensure that relevant meta estimators (e.g. GridSearchCV, Pipeline) are selected
# for patching if they are not already included in the output of `all_estimators()`
estimators_to_patch = set(estimators_to_patch).union(
set(_get_meta_estimators_for_autologging()))
for class_def in estimators_to_patch:
for func_name in ["fit", "fit_transform", "fit_predict"]:
if hasattr(class_def, func_name):
original = getattr(class_def, func_name)
# A couple of estimators use property methods to return fitting functions,
# rather than defining the fitting functions on the estimator class directly.
#
# Example: https://github.com/scikit-learn/scikit-learn/blob/0.23.2/sklearn/neighbors/_lof.py#L183 # noqa
#
# We currently exclude these property fitting methods from patching because
# it's challenging to patch them correctly.
#
# Excluded fitting methods:
# - sklearn.cluster._agglomerative.FeatureAgglomeration.fit_predict
# - sklearn.neighbors._lof.LocalOutlierFactor.fit_predict
#
# You can list property fitting methods by inserting "print(class_def, func_name)"
# in the if clause below.
if isinstance(original, property):
continue
patch_func = create_patch_func(func_name)
wrap_patch(class_def, func_name, patch_func)
def _autolog(log_every_n_epoch=1):
"""
Enable automatic logging from pytorch to MLflow.
Logs loss and any other metrics specified in the fit
function, and optimizer data as parameters. Model checkpoints
are logged as artifacts and pytorch model is stored under `model` directory.
MLflow will also log the parameters of the
`EarlyStoppingCallback <https://pytorch-lightning.readthedocs.io/en/latest/early_stopping.html>`
:param log_every_n_epoch: parameter to log metrics once in `n` epoch. By default, metrics
are logged after every epoch.
"""
global every_n_epoch
every_n_epoch = log_every_n_epoch
class __MLflowPLCallback(pl.Callback):
"""
Callback for auto-logging metrics and parameters.
"""
def __init__(self):
self.early_stopping = False
def on_epoch_end(self, trainer, pl_module):
"""
Log loss and other metrics values after each epoch
:param trainer: pytorch lightning trainer instance
:param pl_module: pytorch lightning base module
"""
if (pl_module.current_epoch + 1) % every_n_epoch == 0:
for key, value in trainer.callback_metrics.items():
try_mlflow_log(mlflow.log_metric,
key,
float(value),
step=pl_module.current_epoch)
for callback in trainer.callbacks:
if isinstance(callback,
pl.callbacks.early_stopping.EarlyStopping):
self._early_stop_check(callback)
def on_train_start(self, trainer, pl_module):
"""
Logs Optimizer related metrics when the train begins
:param trainer: pytorch lightning trainer instance
:param pl_module: pytorch lightning base module
"""
try_mlflow_log(mlflow.set_tag, "Mode", "training")
try_mlflow_log(mlflow.log_param, "epochs", trainer.max_epochs)
for callback in trainer.callbacks:
if isinstance(callback,
pl.callbacks.early_stopping.EarlyStopping):
self.early_stopping = True
self._log_early_stop_params(callback)
# TODO For logging optimizer params - Following scenarios are to revisited.
# 1. In the current scenario, only the first optimizer details are logged.
# Code to be enhanced to log params when multiple optimizers are used.
# 2. mlflow.log_params is used to store optimizer default values into mlflow.
# The keys in default dictionary are too short, Ex: (lr - learning_rate).
# Efficient mapping technique needs to be introduced
# to rename the optimizer parameters based on keys in default dictionary.
if hasattr(trainer, "optimizers"):
optimizer = trainer.optimizers[0]
try_mlflow_log(mlflow.log_param, "optimizer_name",
type(optimizer).__name__)
if hasattr(optimizer, "defaults"):
try_mlflow_log(mlflow.log_params, optimizer.defaults)
summary = str(ModelSummary(pl_module, mode="full"))
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "model_summary.txt")
with open(summary_file, "w") as f:
f.write(summary)
try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def on_train_end(self, trainer, pl_module):
"""
Logs the model checkpoint into mlflow - models folder on the training end
:param trainer: pytorch lightning trainer instance
:param pl_module: pytorch lightning base module
"""
mlflow.pytorch.log_model(pytorch_model=trainer.model,
artifact_path="model")
if self.early_stopping and trainer.checkpoint_callback.best_model_path:
try_mlflow_log(
mlflow.log_artifact,
local_path=trainer.checkpoint_callback.best_model_path,
artifact_path="restored_model_checkpoint",
)
def on_test_end(self, trainer, pl_module):
"""
Logs accuracy and other relevant metrics on the testing end
:param trainer: pytorch lightning trainer instance
:param pl_module: pytorch lightning base module
"""
try_mlflow_log(mlflow.set_tag, "Mode", "testing")
for key, value in trainer.callback_metrics.items():
try_mlflow_log(mlflow.log_metric, key, float(value))
@staticmethod
def _log_early_stop_params(early_stop_obj):
"""
Logs Early Stop parameters into mlflow
:param early_stop_obj: Early stopping callback dict
"""
if hasattr(early_stop_obj, "monitor"):
try_mlflow_log(mlflow.log_param, "monitor",
early_stop_obj.monitor)
if hasattr(early_stop_obj, "mode"):
try_mlflow_log(mlflow.log_param, "mode", early_stop_obj.mode)
if hasattr(early_stop_obj, "patience"):
try_mlflow_log(mlflow.log_param, "patience",
early_stop_obj.patience)
if hasattr(early_stop_obj, "min_delta"):
try_mlflow_log(mlflow.log_param, "min_delta",
early_stop_obj.min_delta)
if hasattr(early_stop_obj, "stopped_epoch"):
try_mlflow_log(mlflow.log_param, "stopped_epoch",
early_stop_obj.stopped_epoch)
@staticmethod
def _early_stop_check(early_stop_callback):
"""
Logs all early stopping metrics
:param early_stop_callback: Early stopping callback object
"""
if early_stop_callback.stopped_epoch != 0:
if hasattr(early_stop_callback, "stopped_epoch"):
try_mlflow_log(mlflow.log_metric, "stopped_epoch",
early_stop_callback.stopped_epoch)
restored_epoch = early_stop_callback.stopped_epoch - max(
1, early_stop_callback.patience)
try_mlflow_log(mlflow.log_metric, "restored_epoch",
restored_epoch)
if hasattr(early_stop_callback, "best_score"):
try_mlflow_log(mlflow.log_metric, "best_score",
float(early_stop_callback.best_score))
if hasattr(early_stop_callback, "wait_count"):
try_mlflow_log(mlflow.log_metric, "wait_count",
early_stop_callback.wait_count)
def _run_and_log_function(self, original, args, kwargs):
"""
This method would be called from patched fit method and
It adds the custom callback class into callback list.
"""
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
if not any(
isinstance(callbacks, __MLflowPLCallback)
for callbacks in self.callbacks):
self.callbacks += [__MLflowPLCallback()]
result = original(self, *args, **kwargs)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return result
@gorilla.patch(pl.Trainer)
def fit(self, *args, **kwargs):
"""
Patching trainer.fit method to add autolog class into callback
"""
original = gorilla.get_original_attribute(pl.Trainer, "fit")
return _run_and_log_function(self, original, args, kwargs)
wrap_patch(pl.Trainer, "fit", fit)
def autolog(importance_types=["weight"]): # pylint: disable=W0102
"""
Enables automatic logging from XGBoost to MLflow. Logs the following.
- parameters specified in `xgboost.train`_.
- metrics on each iteration (if ``evals`` specified).
- metrics at the best iteration (if ``early_stopping_rounds`` specified).
- feature importance as JSON files and plots.
- trained model, including:
- an example of valid input.
- inferred signature of the inputs and outputs of the model.
Note that the `scikit-learn API`_ is not supported.
:param importance_types: importance types to log.
"""
import xgboost
import numpy as np
class _InputExampleInfo:
def __init__(self, input_example=None, error_msg=None):
self.input_example = input_example
self.error_msg = error_msg
# Patching this function so we can get a copy of the data given to DMatrix.__init__
# to use as an input example and for inferring the model signature.
# (there is no way to get the data back from a DMatrix object)
# We store it on the DMatrix object so the train function is able to read it.
def __init__(self, *args, **kwargs):
data = args[0] if len(args) > 0 else kwargs.get("data")
if data is not None:
original = gorilla.get_original_attribute(xgboost.DMatrix,
"__init__")
try:
if isinstance(data, str):
raise Exception("cannot gather example input when " +
"dataset is loaded from a file.")
input_example_info = _InputExampleInfo(
input_example=deepcopy(data[:INPUT_EXAMPLE_SAMPLE_ROWS]))
except Exception as e: # pylint: disable=broad-except
input_example_info = _InputExampleInfo(error_msg=str(e))
setattr(self, "input_example_info", input_example_info)
original(self, *args, **kwargs)
def train(*args, **kwargs):
def record_eval_results(eval_results):
"""
Create a callback function that records evaluation results.
"""
def callback(env):
eval_results.append(dict(env.evaluation_result_list))
return callback
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
def log_feature_importance_plot(features, importance, importance_type):
"""
Log feature importance plot.
"""
import matplotlib.pyplot as plt
features = np.array(features)
importance = np.array(importance)
indices = np.argsort(importance)
features = features[indices]
importance = importance[indices]
num_features = len(features)
# If num_features > 10, increase the figure height to prevent the plot
# from being too dense.
w, h = [6.4, 4.8] # matplotlib's default figure size
h = h + 0.1 * num_features if num_features > 10 else h
fig, ax = plt.subplots(figsize=(w, h))
yloc = np.arange(num_features)
ax.barh(yloc, importance, align="center", height=0.5)
ax.set_yticks(yloc)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
fig.tight_layout()
tmpdir = tempfile.mkdtemp()
try:
# pylint: disable=undefined-loop-variable
filepath = os.path.join(
tmpdir, "feature_importance_{}.png".format(imp_type))
fig.savefig(filepath)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
original = gorilla.get_original_attribute(xgboost, "train")
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
params = args[0] if len(args) > 0 else kwargs["params"]
try_mlflow_log(mlflow.log_params, params)
unlogged_params = [
"params",
"dtrain",
"evals",
"obj",
"feval",
"evals_result",
"xgb_model",
"callbacks",
"learning_rates",
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
all_arg_names = inspect.getargspec(original)[0] # pylint: disable=W1505
num_pos_args = len(args)
# adding a callback that records evaluation results.
eval_results = []
callbacks_index = all_arg_names.index("callbacks")
callback = record_eval_results(eval_results)
if num_pos_args >= callbacks_index + 1:
tmp_list = list(args)
tmp_list[callbacks_index] += [callback]
args = tuple(tmp_list)
elif "callbacks" in kwargs and kwargs["callbacks"] is not None:
kwargs["callbacks"] += [callback]
else:
kwargs["callbacks"] = [callback]
# training model
model = original(*args, **kwargs)
# logging metrics on each iteration.
for idx, metrics in enumerate(eval_results):
try_mlflow_log(mlflow.log_metrics, metrics, step=idx)
# If early_stopping_rounds is present, logging metrics at the best iteration
# as extra metrics with the max step + 1.
early_stopping_index = all_arg_names.index("early_stopping_rounds")
early_stopping = (num_pos_args >= early_stopping_index + 1
or "early_stopping_rounds" in kwargs)
if early_stopping:
extra_step = len(eval_results)
try_mlflow_log(mlflow.log_metric, "stopped_iteration",
len(eval_results) - 1)
try_mlflow_log(mlflow.log_metric, "best_iteration",
model.best_iteration)
try_mlflow_log(mlflow.log_metrics,
eval_results[model.best_iteration],
step=extra_step)
# logging feature importance as artifacts.
for imp_type in importance_types:
imp = None
try:
imp = model.get_score(importance_type=imp_type)
features, importance = zip(*imp.items())
log_feature_importance_plot(features, importance, imp_type)
except Exception: # pylint: disable=broad-except
_logger.exception(
"Failed to log feature importance plot. XGBoost autologging "
"will ignore the failure and continue. Exception: ")
if imp is not None:
tmpdir = tempfile.mkdtemp()
try:
filepath = os.path.join(
tmpdir, "feature_importance_{}.json".format(imp_type))
with open(filepath, "w") as f:
json.dump(imp, f)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
shutil.rmtree(tmpdir)
# dtrain must exist as the original train function already ran successfully
dtrain = args[1] if len(args) > 1 else kwargs.get("dtrain")
input_example = None
signature = None
try:
# it is possible that the dataset was constructed before the patched
# constructor was applied, so we cannot assume the input_example_info exists
input_example_info = getattr(dtrain, "input_example_info", None)
if input_example_info is None:
raise Exception("please ensure that autologging is " +
"enabled before constructing the dataset.")
input_example = input_example_info.input_example
if input_example is None:
# input example collection failed
raise Exception(input_example_info.error_msg)
model_output = model.predict(xgboost.DMatrix(input_example))
signature = infer_signature(input_example, model_output)
except Exception as e: # pylint: disable=broad-except
input_example = None
msg = "Failed to gather example input and model signature: " + str(
e)
_logger.warning(msg)
try_mlflow_log(
log_model,
model,
artifact_path="model",
signature=signature,
input_example=input_example,
)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return model
wrap_patch(xgboost, "train", train)
wrap_patch(xgboost.DMatrix, "__init__", __init__)
def autolog(log_input_example=False, log_model_signature=True):
"""
Enables automatic logging from LightGBM to MLflow. Logs the following.
- parameters specified in `lightgbm.train`_.
- metrics on each iteration (if ``valid_sets`` specified).
- metrics at the best iteration (if ``early_stopping_rounds`` specified).
- feature importance (both "split" and "gain") as JSON files and plots.
- trained model, including:
- an example of valid input.
- inferred signature of the inputs and outputs of the model.
Note that the `scikit-learn API`_ is not supported.
:param log_input_example: if True, logs a sample of the training data as part of the model
as an example for future reference. If False, no sample is logged.
:param log_model_signature: if True, records the type signature of the inputs and outputs as
part of the model. If False, the signature is not recorded to the
model.
"""
import lightgbm
import numpy as np
# Patching this function so we can get a copy of the data given to Dataset.__init__
# to use as an input example and for inferring the model signature.
# (there is no way to get the data back from a Dataset object once it is consumed by train)
# We store it on the Dataset object so the train function is able to read it.
def __init__(self, *args, **kwargs):
data = args[0] if len(args) > 0 else kwargs.get("data")
if data is not None:
original = gorilla.get_original_attribute(lightgbm.Dataset, "__init__")
try:
if isinstance(data, str):
raise Exception(
"cannot gather example input when dataset is loaded from a file."
)
input_example_info = _InputExampleInfo(
input_example=deepcopy(data[:INPUT_EXAMPLE_SAMPLE_ROWS])
)
except Exception as e: # pylint: disable=broad-except
input_example_info = _InputExampleInfo(error_msg=str(e))
setattr(self, "input_example_info", input_example_info)
original(self, *args, **kwargs)
def train(*args, **kwargs):
def record_eval_results(eval_results):
"""
Create a callback function that records evaluation results.
"""
def callback(env):
res = {}
for data_name, eval_name, value, _ in env.evaluation_result_list:
key = data_name + "-" + eval_name
res[key] = value
eval_results.append(res)
return callback
def log_feature_importance_plot(features, importance, importance_type):
"""
Log feature importance plot.
"""
import matplotlib.pyplot as plt
indices = np.argsort(importance)
features = np.array(features)[indices]
importance = importance[indices]
num_features = len(features)
# If num_features > 10, increase the figure height to prevent the plot
# from being too dense.
w, h = [6.4, 4.8] # matplotlib's default figure size
h = h + 0.1 * num_features if num_features > 10 else h
fig, ax = plt.subplots(figsize=(w, h))
yloc = np.arange(num_features)
ax.barh(yloc, importance, align="center", height=0.5)
ax.set_yticks(yloc)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
fig.tight_layout()
tmpdir = tempfile.mkdtemp()
try:
# pylint: disable=undefined-loop-variable
filepath = os.path.join(tmpdir, "feature_importance_{}.png".format(imp_type))
fig.savefig(filepath)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
original = gorilla.get_original_attribute(lightgbm, "train")
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
params = args[0] if len(args) > 0 else kwargs["params"]
try_mlflow_log(mlflow.log_params, params)
unlogged_params = [
"params",
"train_set",
"valid_sets",
"valid_names",
"fobj",
"feval",
"init_model",
"evals_result",
"learning_rates",
"callbacks",
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
all_arg_names = inspect.getargspec(original)[0] # pylint: disable=W1505
num_pos_args = len(args)
# adding a callback that records evaluation results.
eval_results = []
callbacks_index = all_arg_names.index("callbacks")
callback = record_eval_results(eval_results)
if num_pos_args >= callbacks_index + 1:
tmp_list = list(args)
tmp_list[callbacks_index] += [callback]
args = tuple(tmp_list)
elif "callbacks" in kwargs and kwargs["callbacks"] is not None:
kwargs["callbacks"] += [callback]
else:
kwargs["callbacks"] = [callback]
# training model
model = original(*args, **kwargs)
# logging metrics on each iteration.
for idx, metrics in enumerate(eval_results):
try_mlflow_log(mlflow.log_metrics, metrics, step=idx)
# If early_stopping_rounds is present, logging metrics at the best iteration
# as extra metrics with the max step + 1.
early_stopping_index = all_arg_names.index("early_stopping_rounds")
early_stopping = (
num_pos_args >= early_stopping_index + 1 or "early_stopping_rounds" in kwargs
)
if early_stopping:
extra_step = len(eval_results)
try_mlflow_log(mlflow.log_metric, "stopped_iteration", len(eval_results))
# best_iteration is set even if training does not stop early.
try_mlflow_log(mlflow.log_metric, "best_iteration", model.best_iteration)
# iteration starts from 1 in LightGBM.
try_mlflow_log(
mlflow.log_metrics, eval_results[model.best_iteration - 1], step=extra_step
)
# logging feature importance as artifacts.
for imp_type in ["split", "gain"]:
features = model.feature_name()
importance = model.feature_importance(importance_type=imp_type)
try:
log_feature_importance_plot(features, importance, imp_type)
except Exception: # pylint: disable=broad-except
_logger.exception(
"Failed to log feature importance plot. LightGBM autologging "
"will ignore the failure and continue. Exception: "
)
imp = {ft: imp for ft, imp in zip(features, importance.tolist())}
tmpdir = tempfile.mkdtemp()
try:
filepath = os.path.join(tmpdir, "feature_importance_{}.json".format(imp_type))
with open(filepath, "w") as f:
json.dump(imp, f, indent=2)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
shutil.rmtree(tmpdir)
# train_set must exist as the original train function already ran successfully
train_set = args[1] if len(args) > 1 else kwargs.get("train_set")
# it is possible that the dataset was constructed before the patched
# constructor was applied, so we cannot assume the input_example_info exists
input_example_info = getattr(train_set, "input_example_info", None)
def get_input_example():
if input_example_info is None:
raise Exception(ENSURE_AUTOLOGGING_ENABLED_TEXT)
if input_example_info.error_msg is not None:
raise Exception(input_example_info.error_msg)
return input_example_info.input_example
def infer_model_signature(input_example):
model_output = model.predict(input_example)
model_signature = infer_signature(input_example, model_output)
return model_signature
input_example, signature = resolve_input_example_and_signature(
get_input_example,
infer_model_signature,
log_input_example,
log_model_signature,
_logger,
)
try_mlflow_log(
log_model,
model,
artifact_path="model",
signature=signature,
input_example=input_example,
)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return model
wrap_patch(lightgbm, "train", train)
wrap_patch(lightgbm.Dataset, "__init__", __init__)
def autolog(every_n_iter=100):
# pylint: disable=E0611
"""
Enables automatic logging from TensorFlow to MLflow.
Note that autologging for ``tf.keras`` is handled by :py:func:`mlflow.tensorflow.autolog`,
not :py:func:`mlflow.keras.autolog`.
As an example, try running the
`TensorFlow examples <https://github.com/mlflow/mlflow/tree/master/examples/tensorflow>`_.
For each TensorFlow module, autologging captures the following information:
**tf.keras**
- **Metrics** and **Parameters**
- Training loss; validation loss; user-specified metrics
- ``fit()`` or ``fit_generator()`` parameters; optimizer name; learning rate; epsilon
- **Artifacts**
- Model summary on training start
- `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model)
- TensorBoard logs on training end
**tf.keras.callbacks.EarlyStopping**
- **Metrics** and **Parameters**
- Metrics from the ``EarlyStopping`` callbacks: ``stopped_epoch``, ``restored_epoch``,
``restore_best_weight``, etc
- ``fit()`` or ``fit_generator()`` parameters associated with ``EarlyStopping``:
``min_delta``, ``patience``, ``baseline``, ``restore_best_weights``, etc
**tf.estimator**
- **Metrics** and **Parameters**
- TensorBoard metrics: ``average_loss``, ``loss``, etc
- Parameters ``steps`` and ``max_steps``
- **Artifacts**
- `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (TF saved model) on call
to ``tf.estimator.export_saved_model``
**TensorFlow Core**
- **Metrics**
- All ``tf.summary.scalar`` calls
Refer to the autologging tracking documentation for more
information on `TensorFlow workflows
<https://www.mlflow.org/docs/latest/tracking.html#tensorflow-and-keras-experimental>`_.
:param every_n_iter: The frequency with which metrics should be logged.
Defaults to 100. Ex: a value of 100 will log metrics
at step 0, 100, 200, etc.
"""
import tensorflow
global _LOG_EVERY_N_STEPS
_LOG_EVERY_N_STEPS = every_n_iter
atexit.register(_flush_queue)
if LooseVersion(tensorflow.__version__) < LooseVersion("1.12"):
warnings.warn(
"Could not log to MLflow. TensorFlow versions below 1.12 are not supported."
)
return
try:
from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
from tensorflow.python.summary.writer.event_file_writer_v2 import EventFileWriterV2
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.summary.writer.writer import FileWriter
except ImportError:
warnings.warn(
"Could not log to MLflow. TensorFlow versions below 1.12 are not supported."
)
return
@contextmanager
def _manage_active_run():
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
global _AUTOLOG_RUN_ID
if mlflow.active_run(
) is not None: # defensive check in case `mlflow.start_run` fails
_AUTOLOG_RUN_ID = mlflow.active_run().info.run_id
yield mlflow.active_run()
if mlflow.active_run() is not None and mlflow.active_run(
).info.run_id == _AUTOLOG_RUN_ID:
try_mlflow_log(mlflow.end_run)
def train(self, *args, **kwargs):
with _manage_active_run():
original = gorilla.get_original_attribute(
tensorflow.estimator.Estimator, "train")
# Checking step and max_step parameters for logging
if len(args) >= 3:
try_mlflow_log(mlflow.log_param, "steps", args[2])
if len(args) >= 4:
try_mlflow_log(mlflow.log_param, "max_steps", args[3])
if "steps" in kwargs:
try_mlflow_log(mlflow.log_param, "steps", kwargs["steps"])
if "max_steps" in kwargs:
try_mlflow_log(mlflow.log_param, "max_steps",
kwargs["max_steps"])
result = original(self, *args, **kwargs)
return result
def export_saved_model(self, *args, **kwargs):
auto_end = False
if not mlflow.active_run():
global _AUTOLOG_RUN_ID
if _AUTOLOG_RUN_ID:
try_mlflow_log(mlflow.start_run, _AUTOLOG_RUN_ID)
else:
try_mlflow_log(mlflow.start_run)
auto_end = True
original = gorilla.get_original_attribute(
tensorflow.estimator.Estimator, "export_saved_model")
serialized = original(self, *args, **kwargs)
try_mlflow_log(
log_model,
tf_saved_model_dir=serialized.decode("utf-8"),
tf_meta_graph_tags=[tag_constants.SERVING],
tf_signature_def_key="predict",
artifact_path="model",
)
if (mlflow.active_run() is not None and mlflow.active_run().info.run_id
== _AUTOLOG_RUN_ID) or auto_end:
try_mlflow_log(mlflow.end_run)
return serialized
def export_savedmodel(self, *args, **kwargs):
auto_end = False
global _AUTOLOG_RUN_ID
if not mlflow.active_run():
if _AUTOLOG_RUN_ID:
try_mlflow_log(mlflow.start_run, _AUTOLOG_RUN_ID)
else:
try_mlflow_log(mlflow.start_run)
auto_end = True
original = gorilla.get_original_attribute(
tensorflow.estimator.Estimator, "export_savedmodel")
serialized = original(self, *args, **kwargs)
try_mlflow_log(
log_model,
tf_saved_model_dir=serialized.decode("utf-8"),
tf_meta_graph_tags=[tag_constants.SERVING],
tf_signature_def_key="predict",
artifact_path="model",
)
if (mlflow.active_run() is not None and mlflow.active_run().info.run_id
== _AUTOLOG_RUN_ID) or auto_end:
try_mlflow_log(mlflow.end_run)
return serialized
def _early_stop_check(callbacks):
for callback in callbacks:
if isinstance(callback, tensorflow.keras.callbacks.EarlyStopping):
return callback
return None
def _log_early_stop_callback_params(callback):
if callback:
try:
earlystopping_params = {
"monitor": callback.monitor,
"min_delta": callback.min_delta,
"patience": callback.patience,
"baseline": callback.baseline,
"restore_best_weights": callback.restore_best_weights,
}
try_mlflow_log(mlflow.log_params, earlystopping_params)
except Exception: # pylint: disable=W0703
return
def _get_early_stop_callback_attrs(callback):
try:
return callback.stopped_epoch, callback.restore_best_weights, callback.patience
except Exception: # pylint: disable=W0703
return None
def _log_early_stop_callback_metrics(callback, history):
if callback:
callback_attrs = _get_early_stop_callback_attrs(callback)
if callback_attrs is None:
return
stopped_epoch, restore_best_weights, patience = callback_attrs
try_mlflow_log(mlflow.log_metric, "stopped_epoch", stopped_epoch)
# Weights are restored only if early stopping occurs
if stopped_epoch != 0 and restore_best_weights:
restored_epoch = stopped_epoch - max(1, patience)
try_mlflow_log(mlflow.log_metric, "restored_epoch",
restored_epoch)
restored_metrics = {
key: history.history[key][restored_epoch]
for key in history.history.keys()
}
# Metrics are logged as 'epoch_loss' and 'epoch_acc' in TF 1.X
if LooseVersion(
tensorflow.__version__) < LooseVersion("2.0.0"):
if "loss" in restored_metrics:
restored_metrics["epoch_loss"] = restored_metrics.pop(
"loss")
if "acc" in restored_metrics:
restored_metrics["epoch_acc"] = restored_metrics.pop(
"acc")
# Checking that a metric history exists
metric_key = next(iter(history.history), None)
if metric_key is not None:
last_epoch = len(history.history[metric_key])
try_mlflow_log(mlflow.log_metrics,
restored_metrics,
step=last_epoch)
def fit(self, *args, **kwargs):
with _manage_active_run():
original = gorilla.get_original_attribute(tensorflow.keras.Model,
"fit")
unlogged_params = [
"self", "x", "y", "callbacks", "validation_data", "verbose"
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
early_stop_callback = None
# Checking if the 'callback' argument of fit() is set
if len(args) >= 6:
tmp_list = list(args)
early_stop_callback = _early_stop_check(tmp_list[5])
tmp_list[5], log_dir = _setup_callbacks(tmp_list[5])
args = tuple(tmp_list)
elif "callbacks" in kwargs:
early_stop_callback = _early_stop_check(kwargs["callbacks"])
kwargs["callbacks"], log_dir = _setup_callbacks(
kwargs["callbacks"])
else:
kwargs["callbacks"], log_dir = _setup_callbacks([])
_log_early_stop_callback_params(early_stop_callback)
history = original(self, *args, **kwargs)
_log_early_stop_callback_metrics(early_stop_callback, history)
_flush_queue()
_log_artifacts_with_warning(local_dir=log_dir.location,
artifact_path="tensorboard_logs")
if log_dir.is_temp:
shutil.rmtree(log_dir.location)
return history
def fit_generator(self, *args, **kwargs):
with _manage_active_run():
original = gorilla.get_original_attribute(tensorflow.keras.Model,
"fit_generator")
unlogged_params = [
"self", "generator", "callbacks", "validation_data", "verbose"
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
# Checking if the 'callback' argument of fit() is set
if len(args) >= 5:
tmp_list = list(args)
tmp_list[4], log_dir = _setup_callbacks(tmp_list[4])
args = tuple(tmp_list)
elif "callbacks" in kwargs:
kwargs["callbacks"], log_dir = _setup_callbacks(
kwargs["callbacks"])
else:
kwargs["callbacks"], log_dir = _setup_callbacks([])
result = original(self, *args, **kwargs)
_flush_queue()
_log_artifacts_with_warning(local_dir=log_dir.location,
artifact_path="tensorboard_logs")
if log_dir.is_temp:
shutil.rmtree(log_dir.location)
return result
def add_event(self, event):
_log_event(event)
original = gorilla.get_original_attribute(EventFileWriter, "add_event")
return original(self, event)
def add_summary(self, *args, **kwargs):
original = gorilla.get_original_attribute(FileWriter, "add_summary")
result = original(self, *args, **kwargs)
_flush_queue()
return result
patches = [
(EventFileWriter, "add_event", add_event),
(EventFileWriterV2, "add_event", add_event),
(tensorflow.estimator.Estimator, "train", train),
(tensorflow.keras.Model, "fit", fit),
(tensorflow.keras.Model, "fit_generator", fit_generator),
(tensorflow.estimator.Estimator, "export_saved_model",
export_saved_model),
(tensorflow.estimator.Estimator, "export_savedmodel",
export_savedmodel),
(FileWriter, "add_summary", add_summary),
]
for p in patches:
wrap_patch(*p)
def autolog():
"""
Enable automatic logging from Fastai to MLflow.
Logs loss and any other metrics specified in the fit
function, and optimizer data as parameters. Model checkpoints
are logged as artifacts to a 'models' directory.
MLflow will also log the parameters of the
`EarlyStoppingCallback <https://docs.fast.ai/callbacks.html#EarlyStoppingCallback>`_
and `OneCycleScheduler <https://docs.fast.ai/callbacks.html#OneCycleScheduler>`_ callbacks
"""
from fastai.basic_train import LearnerCallback, Learner
from fastai.callbacks.hooks import model_summary, layers_info
from fastai.callbacks import EarlyStoppingCallback, OneCycleScheduler
class __MLflowFastaiCallback(LearnerCallback):
"""
Callback for auto-logging metrics and parameters.
Records model structural information as params when training begins
"""
def __init__(
self,
learner,
):
super().__init__(learner)
self.learner = learner
self.opt = self.learn.opt
self.metrics_names = ["train_loss", "valid_loss"
] + [o.__name__ for o in learner.metrics]
def on_epoch_end(self, **kwargs):
"""
Log loss and other metrics values after each epoch
"""
if kwargs["smooth_loss"] is None or kwargs["last_metrics"] is None:
return
epoch = kwargs["epoch"]
metrics = [kwargs["smooth_loss"]] + kwargs["last_metrics"]
metrics = map(float, metrics)
metrics = dict(zip(self.metrics_names, metrics))
try_mlflow_log(mlflow.log_metrics, metrics, step=epoch)
def on_train_begin(self, **kwargs):
info = layers_info(self.learner)
try_mlflow_log(mlflow.log_param, "num_layers", len(info))
try_mlflow_log(mlflow.log_param, "opt_func",
self.opt_func.func.__name__)
if hasattr(self.opt, "true_wd"):
try_mlflow_log(mlflow.log_param, "true_wd", self.opt.true_wd)
if hasattr(self.opt, "bn_wd"):
try_mlflow_log(mlflow.log_param, "bn_wd", self.opt.bn_wd)
if hasattr(self.opt, "train_bn"):
try_mlflow_log(mlflow.log_param, "train_bn", self.train_bn)
summary = model_summary(self.learner)
try_mlflow_log(mlflow.set_tag, "model_summary", summary)
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "model_summary.txt")
with open(summary_file, "w") as f:
f.write(summary)
try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def on_train_end(self, **kwargs):
try_mlflow_log(log_model, self.learner, artifact_path="model")
def _find_callback_of_type(callback_type, callbacks):
for callback in callbacks:
if isinstance(callback, callback_type):
return callback
return None
def _log_early_stop_callback_params(callback):
if callback:
try:
earlystopping_params = {
"early_stop_monitor": callback.monitor,
"early_stop_min_delta": callback.min_delta,
"early_stop_patience": callback.patience,
"early_stop_mode": callback.mode,
}
try_mlflow_log(mlflow.log_params, earlystopping_params)
except Exception: # pylint: disable=W0703
return
def _log_one_cycle_callback_params(callback):
if callback:
try:
params = {
"lr_max": callback.lr_max,
"div_factor": callback.div_factor,
"pct_start": callback.pct_start,
"final_div": callback.final_div,
"tot_epochs": callback.tot_epochs,
"start_epoch": callback.start_epoch,
"moms": callback.moms,
}
try_mlflow_log(mlflow.log_params, params)
except Exception: # pylint: disable=W0703
return
def _run_and_log_function(self, original, args, kwargs, unlogged_params,
callback_arg_index):
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
log_fn_args_as_params(original, [self] + list(args), kwargs,
unlogged_params)
callbacks = [cb(self)
for cb in self.callback_fns] + (self.callbacks or [])
# Checking if the 'callback' argument of the function is set
if len(args) > callback_arg_index:
tmp_list = list(args)
callbacks += list(args[callback_arg_index])
tmp_list[callback_arg_index] += [__MLflowFastaiCallback(self)]
args = tuple(tmp_list)
elif "callbacks" in kwargs:
callbacks += list(kwargs["callbacks"])
kwargs["callbacks"] += [__MLflowFastaiCallback(self)]
else:
kwargs["callbacks"] = [__MLflowFastaiCallback(self)]
early_stop_callback = _find_callback_of_type(EarlyStoppingCallback,
callbacks)
one_cycle_callback = _find_callback_of_type(OneCycleScheduler,
callbacks)
_log_early_stop_callback_params(early_stop_callback)
_log_one_cycle_callback_params(one_cycle_callback)
result = original(self, *args, **kwargs)
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return result
def fit(self, *args, **kwargs):
original = gorilla.get_original_attribute(Learner, "fit")
unlogged_params = ["self", "callbacks", "learner"]
return _run_and_log_function(self, original, args, kwargs,
unlogged_params, 3)
wrap_patch(Learner, "fit", fit)
def autolog(importance_types=["weight"]): # pylint: disable=W0102
"""
Enables automatic logging from XGBoost to MLflow. Logs the following.
- parameters specified in `xgboost.train`_.
- metrics on each iteration (if ``evals`` specified).
- metrics at the best iteration (if ``early_stopping_rounds`` specified).
- feature importance as JSON files and plots.
- trained model.
Note that the `scikit-learn API`_ is not supported.
:param importance_types: importance types to log.
"""
import xgboost
import numpy as np
def train(*args, **kwargs):
def record_eval_results(eval_results):
"""
Create a callback function that records evaluation results.
"""
def callback(env):
eval_results.append(dict(env.evaluation_result_list))
return callback
if not mlflow.active_run():
try_mlflow_log(mlflow.start_run)
auto_end_run = True
else:
auto_end_run = False
def log_feature_importance_plot(features, importance, importance_type):
"""
Log feature importance plot.
"""
import matplotlib.pyplot as plt
features = np.array(features)
importance = np.array(importance)
indices = np.argsort(importance)
features = features[indices]
importance = importance[indices]
num_features = len(features)
# If num_features > 10, increase the figure height to prevent the plot
# from being too dense.
w, h = [6.4, 4.8] # matplotlib's default figure size
h = h + 0.1 * num_features if num_features > 10 else h
fig, ax = plt.subplots(figsize=(w, h))
yloc = np.arange(num_features)
ax.barh(yloc, importance, align="center", height=0.5)
ax.set_yticks(yloc)
ax.set_yticklabels(features)
ax.set_xlabel("Importance")
ax.set_title("Feature Importance ({})".format(importance_type))
fig.tight_layout()
tmpdir = tempfile.mkdtemp()
try:
# pylint: disable=undefined-loop-variable
filepath = os.path.join(tmpdir, "feature_importance_{}.png".format(imp_type))
fig.savefig(filepath)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
plt.close(fig)
shutil.rmtree(tmpdir)
original = gorilla.get_original_attribute(xgboost, "train")
# logging booster params separately via mlflow.log_params to extract key/value pairs
# and make it easier to compare them across runs.
params = args[0] if len(args) > 0 else kwargs["params"]
try_mlflow_log(mlflow.log_params, params)
unlogged_params = [
"params",
"dtrain",
"evals",
"obj",
"feval",
"evals_result",
"xgb_model",
"callbacks",
"learning_rates",
]
log_fn_args_as_params(original, args, kwargs, unlogged_params)
all_arg_names = inspect.getargspec(original)[0] # pylint: disable=W1505
num_pos_args = len(args)
# adding a callback that records evaluation results.
eval_results = []
callbacks_index = all_arg_names.index("callbacks")
callback = record_eval_results(eval_results)
if num_pos_args >= callbacks_index + 1:
tmp_list = list(args)
tmp_list[callbacks_index] += [callback]
args = tuple(tmp_list)
elif "callbacks" in kwargs and kwargs["callbacks"] is not None:
kwargs["callbacks"] += [callback]
else:
kwargs["callbacks"] = [callback]
# training model
model = original(*args, **kwargs)
# logging metrics on each iteration.
for idx, metrics in enumerate(eval_results):
try_mlflow_log(mlflow.log_metrics, metrics, step=idx)
# If early_stopping_rounds is present, logging metrics at the best iteration
# as extra metrics with the max step + 1.
early_stopping_index = all_arg_names.index("early_stopping_rounds")
early_stopping = (
num_pos_args >= early_stopping_index + 1 or "early_stopping_rounds" in kwargs
)
if early_stopping:
extra_step = len(eval_results)
try_mlflow_log(mlflow.log_metric, "stopped_iteration", len(eval_results) - 1)
try_mlflow_log(mlflow.log_metric, "best_iteration", model.best_iteration)
try_mlflow_log(mlflow.log_metrics, eval_results[model.best_iteration], step=extra_step)
# logging feature importance as artifacts.
for imp_type in importance_types:
imp = None
try:
imp = model.get_score(importance_type=imp_type)
features, importance = zip(*imp.items())
log_feature_importance_plot(features, importance, imp_type)
except Exception: # pylint: disable=broad-except
_logger.exception(
"Failed to log feature importance plot. XGBoost autologging "
"will ignore the failure and continue. Exception: "
)
if imp is not None:
tmpdir = tempfile.mkdtemp()
try:
filepath = os.path.join(tmpdir, "feature_importance_{}.json".format(imp_type))
with open(filepath, "w") as f:
json.dump(imp, f)
try_mlflow_log(mlflow.log_artifact, filepath)
finally:
shutil.rmtree(tmpdir)
try_mlflow_log(log_model, model, artifact_path="model")
if auto_end_run:
try_mlflow_log(mlflow.end_run)
return model
wrap_patch(xgboost, "train", train)
