75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
import logging
|
|
|
|
import xgboost
|
|
from packaging.version import Version
|
|
|
|
# Suppress a false positive pylint error: https://github.com/PyCQA/pylint/issues/1630
|
|
from mlflow.utils.autologging_utils import ExceptionSafeAbstractClass
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _patch_metric_names(metric_dict):
|
|
# XGBoost provides some metrics with "@", e.g. "ndcg@3" that are not valid MLflow metric names
|
|
patched_metrics = {
|
|
metric_name.replace("@", "_at_"): value for metric_name, value in metric_dict.items()
|
|
}
|
|
changed_keys = set(patched_metrics.keys()) - set(metric_dict.keys())
|
|
if changed_keys:
|
|
_logger.info(
|
|
"Identified one or more metrics with names containing the invalid character `@`."
|
|
" These metric names have been sanitized by replacing `@` with `_at_`, as follows: %s",
|
|
", ".join(changed_keys),
|
|
)
|
|
|
|
return patched_metrics
|
|
|
|
|
|
def autolog_callback(env, metrics_logger, eval_results):
|
|
metric_dict = _patch_metric_names(dict(env.evaluation_result_list))
|
|
metrics_logger.record_metrics(metric_dict, env.iteration)
|
|
eval_results.append(metric_dict)
|
|
|
|
|
|
IS_TRAINING_CALLBACK_SUPPORTED = Version(xgboost.__version__.replace("SNAPSHOT", "dev")) >= Version(
|
|
"1.3.0"
|
|
)
|
|
|
|
if IS_TRAINING_CALLBACK_SUPPORTED:
|
|
|
|
class AutologCallback(
|
|
xgboost.callback.TrainingCallback,
|
|
metaclass=ExceptionSafeAbstractClass,
|
|
):
|
|
def __init__(self, metrics_logger, eval_results):
|
|
self.metrics_logger = metrics_logger
|
|
self.eval_results = eval_results
|
|
|
|
def after_iteration(self, model, epoch, evals_log):
|
|
"""
|
|
Run after each iteration. Return True when training should stop.
|
|
"""
|
|
# `evals_log` is a nested dict (type: Dict[str, Dict[str, List[float]]])
|
|
# that looks like this:
|
|
# {
|
|
# "train": {
|
|
# "auc": [0.5, 0.6, 0.7, ...],
|
|
# ...
|
|
# },
|
|
# ...
|
|
# }
|
|
evaluation_result_dict = {}
|
|
for data_name, metric_dict in evals_log.items():
|
|
metric_dict = _patch_metric_names(metric_dict)
|
|
for metric_name, metric_values_on_each_iter in metric_dict.items():
|
|
key = f"{data_name}-{metric_name}"
|
|
# The last element in `metric_values_on_each_iter` corresponds to
|
|
# the metric on the current iteration
|
|
evaluation_result_dict[key] = metric_values_on_each_iter[-1]
|
|
|
|
self.metrics_logger.record_metrics(evaluation_result_dict, epoch)
|
|
self.eval_results.append(evaluation_result_dict)
|
|
|
|
# Return `False` to indicate training should not stop
|
|
return False
|