def _create_child_runs_for_parameter_search(parent_estimator, parent_model,
parent_run, child_tags):
from itertools import zip_longest
client = MlflowClient()
# Use the start time of the parent parameter search run as a rough estimate for the
# start time of child runs, since we cannot precisely determine when each point
# in the parameter search space was explored
child_run_start_time = parent_run.info.start_time
child_run_end_time = int(time.time() * 1000)
estimator_param_maps = parent_estimator.getEstimatorParamMaps()
tuned_estimator = parent_estimator.getEstimator()
metrics_dict, _ = _get_param_search_metrics_and_best_index(
parent_estimator, parent_model)
for i in range(len(estimator_param_maps)):
child_estimator = tuned_estimator.copy(estimator_param_maps[i])
tags_to_log = dict(child_tags) if child_tags else {}
tags_to_log.update({MLFLOW_PARENT_RUN_ID: parent_run.info.run_id})
tags_to_log.update(_get_estimator_info_tags(child_estimator))
child_run = client.create_run(
experiment_id=parent_run.info.experiment_id,
start_time=child_run_start_time,
tags=tags_to_log,
)
params_to_log = _get_instance_param_map(
child_estimator,
parent_estimator._autologging_metadata.uid_to_indexed_name_map)
param_batches_to_log = _chunk_dict(
params_to_log, chunk_size=MAX_PARAMS_TAGS_PER_BATCH)
metrics_to_log = {k: v[i] for k, v in metrics_dict.items()}
for params_batch, metrics_batch in zip_longest(param_batches_to_log,
[metrics_to_log],
fillvalue={}):
# Trim any parameter keys / values and metric keys that exceed the limits
# imposed by corresponding MLflow Tracking APIs (e.g., LogParam, LogMetric)
truncated_params_batch = _truncate_dict(params_batch,
MAX_ENTITY_KEY_LENGTH,
MAX_PARAM_VAL_LENGTH)
truncated_metrics_batch = _truncate_dict(
metrics_batch, max_key_length=MAX_ENTITY_KEY_LENGTH)
client.log_batch(
run_id=child_run.info.run_id,
params=[
Param(str(key), str(value))
for key, value in truncated_params_batch.items()
],
metrics=[
Metric(key=str(key),
value=value,
timestamp=child_run_end_time,
step=0)
for key, value in truncated_metrics_batch.items()
],
)
client.set_terminated(run_id=child_run.info.run_id,
end_time=child_run_end_time)
def _create_child_runs_for_parameter_search(cv_estimator,
parent_run,
child_tags=None):
"""
Creates a collection of child runs for a parameter search training session.
Runs are reconstructed from the `cv_results_` attribute of the specified trained
parameter search estimator - `cv_estimator`, which provides relevant performance
metrics for each point in the parameter search space. One child run is created
for each point in the parameter search space. For additional information, see
`https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.GridSearchCV.html`_. # noqa: E501
:param cv_estimator: The trained parameter search estimator for which to create
child runs.
:param parent_run: A py:class:`mlflow.entities.Run` object referring to the parent
parameter search run for which child runs should be created.
:param child_tags: An optional dictionary of MLflow tag keys and values to log
for each child run.
"""
import pandas as pd
client = MlflowClient()
# Use the start time of the parent parameter search run as a rough estimate for the
# start time of child runs, since we cannot precisely determine when each point
# in the parameter search space was explored
child_run_start_time = parent_run.info.start_time
child_run_end_time = int(time.time() * 1000)
seed_estimator = cv_estimator.estimator
# In the unlikely case that a seed of a parameter search estimator is,
# itself, a parameter search estimator, we should avoid logging the untuned
# parameters of the seeds's seed estimator
should_log_params_deeply = not _is_parameter_search_estimator(
seed_estimator)
# Each row of `cv_results_` only provides parameters that vary across
# the user-specified parameter grid. In order to log the complete set
# of parameters for each child run, we fetch the parameters defined by
# the seed estimator and update them with parameter subset specified
# in the result row
base_params = seed_estimator.get_params(deep=should_log_params_deeply)
cv_results_df = pd.DataFrame.from_dict(cv_estimator.cv_results_)
for _, result_row in cv_results_df.iterrows():
tags_to_log = dict(child_tags) if child_tags else {}
tags_to_log.update({MLFLOW_PARENT_RUN_ID: parent_run.info.run_id})
tags_to_log.update(_get_estimator_info_tags(seed_estimator))
child_run = client.create_run(
experiment_id=parent_run.info.experiment_id,
start_time=child_run_start_time,
tags=tags_to_log,
)
from itertools import zip_longest
params_to_log = dict(base_params)
params_to_log.update(result_row.get("params", {}))
param_batches_to_log = _chunk_dict(
params_to_log, chunk_size=MAX_PARAMS_TAGS_PER_BATCH)
# Parameters values are recorded twice in the set of search `cv_results_`:
# once within a `params` column with dictionary values and once within
# a separate dataframe column that is created for each parameter. To prevent
# duplication of parameters, we log the consolidated values from the parameter
# dictionary column and filter out the other parameter-specific columns with
# names of the form `param_{param_name}`. Additionally, `cv_results_` produces
# metrics for each training split, which is fairly verbose; accordingly, we filter
# out per-split metrics in favor of aggregate metrics (mean, std, etc.)
excluded_metric_prefixes = ["param", "split"]
metric_batches_to_log = _chunk_dict(
{
key: value
for key, value in result_row.iteritems() if not any([
key.startswith(prefix)
for prefix in excluded_metric_prefixes
]) and isinstance(value, Number)
},
chunk_size=min(MAX_ENTITIES_PER_BATCH - MAX_PARAMS_TAGS_PER_BATCH,
MAX_METRICS_PER_BATCH),
)
for params_batch, metrics_batch in zip_longest(param_batches_to_log,
metric_batches_to_log,
fillvalue={}):
# Trim any parameter keys / values and metric keys that exceed the limits
# imposed by corresponding MLflow Tracking APIs (e.g., LogParam, LogMetric)
truncated_params_batch = _truncate_dict(params_batch,
MAX_ENTITY_KEY_LENGTH,
MAX_PARAM_VAL_LENGTH)
truncated_metrics_batch = _truncate_dict(
metrics_batch, max_key_length=MAX_ENTITY_KEY_LENGTH)
client.log_batch(
run_id=child_run.info.run_id,
params=[
Param(str(key), str(value))
for key, value in truncated_params_batch.items()
],
metrics=[
Metric(key=str(key),
value=value,
timestamp=child_run_end_time,
step=0)
for key, value in truncated_metrics_batch.items()
],
)
client.set_terminated(run_id=child_run.info.run_id,
end_time=child_run_end_time)
class MLflowWriter(object):
def __init__(self, exp_name, save_dir, log_every, **mlflow_cfg):
mlflow.set_tracking_uri(save_dir)
self.client = MlflowClient(**mlflow_cfg)
mlflow.set_experiment(exp_name)
self.experiment_id = self.client.get_experiment_by_name(
exp_name).experiment_id
self.run_id = self.client.create_run(self.experiment_id).info.run_id
self.log_every = log_every
self.clear()
def log_params_from_omegaconf(self, params):
self._explore_recursive("", params)
def _explore_recursive(self, parent_name, element):
if isinstance(element, DictConfig):
iterator = element.items()
elif isinstance(element, ListConfig):
iterator = enumerate(element)
for k, v in iterator:
if isinstance(v, DictConfig) or isinstance(v, ListConfig):
self._explore_recursive(f"{parent_name}{k}.", v)
else:
self.client.log_param(
self.run_id, f"{parent_name}{k}", v)
def log_torch_model(self, model, epoch):
with mlflow.start_run(self.run_id):
mlflow.pytorch.log_model(model, "model_%04d" % epoch)
def log_metric(self, key, value, is_training):
if isinstance(value, torch.Tensor):
value = float(value.detach().cpu().numpy())
metric_name = "train/" if is_training else "valid/"
metric_name += str(key)
if key in self.metrics:
self.metrics[metric_name].append(value)
else:
self.metrics[metric_name] = [value]
def next_iteration(self):
self.iterations += 1
if self.iterations % self.log_every == 0:
self.toMlflow(nb_data=self.log_every)
def toMlflow(self, nb_data=0, step=0):
for key, value in self.metrics.items():
self.client.log_metric(
self.run_id, key,
np.mean(value[-nb_data:]), step=step)
def get_mean(self, key, is_training):
metric_name = "train/" if is_training else "valid/"
metric_name += str(key)
return np.mean(self.metrics[metric_name])
def clear(self):
self.metrics = {}
self.iterations = 0
def log_artifact(self, path):
self.client.log_artifact(self.run_id, local_path=path)
def terminate(self):
self.client.set_terminated(self.run_id)
