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_generator(self, *args, **kwargs): 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(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] += [__MLflowKerasCallback()] args = tuple(tmp_list) elif 'callbacks' in kwargs: kwargs['callbacks'] += [__MLflowKerasCallback()] else: kwargs['callbacks'] = [__MLflowKerasCallback()] result = original(self, *args, **kwargs) if auto_end_run: try_mlflow_log(mlflow.end_run) return result
def _patch_implementation( self, original, inst, *args, **kwargs ): # pylint: disable=arguments-differ unlogged_params = ["self", "generator", "callbacks", "validation_data", "verbose"] log_fn_args_as_params(original, args, kwargs, unlogged_params) run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: # Checking if the 'callback' argument of fit() is set if len(args) >= 5: tmp_list = list(args) tmp_list[4], self.log_dir = _setup_callbacks( tmp_list[4], log_models, metrics_logger ) args = tuple(tmp_list) elif kwargs.get("callbacks"): kwargs["callbacks"], self.log_dir = _setup_callbacks( kwargs["callbacks"], log_models, metrics_logger ) else: kwargs["callbacks"], self.log_dir = _setup_callbacks( [], log_models, metrics_logger ) result = original(inst, *args, **kwargs) _flush_queue() _log_artifacts_with_warning( local_dir=self.log_dir.location, artifact_path="tensorboard_logs" ) if self.log_dir.is_temp: shutil.rmtree(self.log_dir.location) return result
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], log_models) args = tuple(tmp_list) elif kwargs.get("callbacks"): early_stop_callback = _early_stop_check(kwargs["callbacks"]) kwargs["callbacks"], log_dir = _setup_callbacks(kwargs["callbacks"], log_models) else: kwargs["callbacks"], log_dir = _setup_callbacks([], log_models) _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 _run_and_log_function(self, original, args, kwargs, unlogged_params, callback_arg_index): log_fn_args_as_params(original, list(args), kwargs, unlogged_params) callbacks = [cb(self) for cb in self.callback_fns] + (self.callbacks or []) run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: mlflowFastaiCallback = getFastaiCallback(metrics_logger, self) # 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] args = tuple(tmp_list) elif kwargs.get("callbacks"): callbacks += list(kwargs["callbacks"]) kwargs["callbacks"] += [mlflowFastaiCallback] else: kwargs["callbacks"] = [mlflowFastaiCallback] 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) return result
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 _run_and_log_function(self, original, args, kwargs, unlogged_params, callback_arg_index): log_fn_args_as_params(original, args, kwargs, unlogged_params) early_stop_callback = None # Checking if the 'callback' argument of the function is set run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: mlflowKerasCallback = getKerasCallback(metrics_logger) 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 kwargs.get("callbacks"): early_stop_callback = _early_stop_check(kwargs["callbacks"]) kwargs["callbacks"] += [mlflowKerasCallback] else: kwargs["callbacks"] = [mlflowKerasCallback] try_mlflow_log(_log_early_stop_callback_params, early_stop_callback) history = original(self, *args, **kwargs) try_mlflow_log(_log_early_stop_callback_metrics, early_stop_callback, history, metrics_logger) return history
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 wrapper_fit(original, self, *args, **kwargs): should_autolog = False if AutologHelpers.should_autolog: AutologHelpers.should_autolog = False should_autolog = True try: if should_autolog: # This may generate warnings due to collisions in already-logged param names log_fn_args_as_params(original, args, kwargs) # training model model = original(self, *args, **kwargs) if should_autolog: # Log the model if get_autologging_config(FLAVOR_NAME, "log_models", True): try_mlflow_log(log_model, model, artifact_path="model") # Log the most common metrics if isinstance(model, statsmodels.base.wrapper.ResultsWrapper): metrics_dict = results_to_dict(model) try_mlflow_log(mlflow.log_metrics, metrics_dict) return model finally: # Clean the shared flag for future calls in case it had been set here ... if should_autolog: AutologHelpers.should_autolog = True
def test_log_fn_args_as_params(args, kwargs, expected, start_run): # pylint: disable=W0613 log_fn_args_as_params(dummy_fn, args, kwargs) client = mlflow.tracking.MlflowClient() params = client.get_run(mlflow.active_run().info.run_id).data.params for arg, value in zip(["arg1", "arg2", "arg3"], expected): assert arg in params assert params[arg] == value
def _run_and_log_function(self, original, args, kwargs, unlogged_params, is_fine_tune=False): # Check if is trying to fit while fine tuning or not mlflow_cbs = [cb for cb in self.cbs if cb.name == "___mlflow_fastai"] fit_in_fine_tune = ( original.__name__ == "fit" and len(mlflow_cbs) > 0 and mlflow_cbs[0].is_fine_tune ) if not fit_in_fine_tune: log_fn_args_as_params(original, list(args), kwargs, unlogged_params) run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: if not fit_in_fine_tune: early_stop_callback = _find_callback_of_type(EarlyStoppingCallback, self.cbs) _log_early_stop_callback_params(early_stop_callback) # First try to remove if any already registered callback self.remove_cbs(mlflow_cbs) # Log information regarding model and data without bar and print-out with self.no_bar(), self.no_logging(): _log_model_info(learner=self) mlflowFastaiCallback = getFastaiCallback( metrics_logger=metrics_logger, is_fine_tune=is_fine_tune ) # Add the new callback self.add_cb(mlflowFastaiCallback) result = original(self, *args, **kwargs) return result
def fit_generator(self, *args, **kwargs): """ NOTE: `fit_generator()` is deprecated in TF >= 2.1.0 and simply wraps `fit()`. To avoid unintentional creation of nested MLflow runs caused by a patched `fit_generator()` method calling a patched `fit()` method, we only patch `fit_generator()` in TF < 2.1.0. """ 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], log_models) args = tuple(tmp_list) elif kwargs.get("callbacks"): kwargs["callbacks"], log_dir = _setup_callbacks(kwargs["callbacks"], log_models) else: kwargs["callbacks"], log_dir = _setup_callbacks([], log_models) 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 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, artifact_path='tensorboard_logs') shutil.rmtree(log_dir) return result
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 test_log_fn_args_as_params_ignores_unwanted_parameters(start_run): # pylint: disable=W0613 args, kwargs, unlogged = ("arg1", { "arg2": "value" }, ["arg1", "arg2", "arg3"]) log_fn_args_as_params(dummy_fn, args, kwargs, unlogged) client = mlflow.tracking.MlflowClient() params = client.get_run(mlflow.active_run().info.run_id).data.params assert len(params.keys()) == 0
def _patch_implementation( self, original, inst, *args, **kwargs ): # pylint: disable=arguments-differ unlogged_params = ["self", "x", "y", "callbacks", "validation_data", "verbose"] log_fn_args_as_params(original, args, kwargs, unlogged_params) early_stop_callback = None run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: # Check if the 'callback' argument of fit() is set positionally if len(args) >= 6: # Convert the positional training function arguments to a list in order to # mutate the contents args = list(args) # Make a shallow copy of the preexisting callbacks to avoid permanently # modifying their contents for future training invocations. Introduce # TensorBoard & tf.keras callbacks if necessary callbacks = list(args[5]) callbacks, self.log_dir = _setup_callbacks( callbacks, log_models, metrics_logger ) # Replace the callbacks positional entry in the copied arguments and convert # the arguments back to tuple form for usage in the training function args[5] = callbacks args = tuple(args) else: # Make a shallow copy of the preexisting callbacks and introduce TensorBoard # & tf.keras callbacks if necessary callbacks = list(kwargs.get("callbacks") or []) kwargs["callbacks"], self.log_dir = _setup_callbacks( callbacks, log_models, metrics_logger ) early_stop_callback = _get_early_stop_callback(callbacks) _log_early_stop_callback_params(early_stop_callback) history = original(inst, *args, **kwargs) _log_early_stop_callback_metrics( callback=early_stop_callback, history=history, metrics_logger=metrics_logger, ) _flush_queue() mlflow.log_artifacts( local_dir=self.log_dir.location, artifact_path="tensorboard_logs", ) if self.log_dir.is_temp: shutil.rmtree(self.log_dir.location) return history
def wrapper_fit(original, self, *args, **kwargs): should_autolog = False if AutologHelpers.should_autolog: AutologHelpers.should_autolog = False should_autolog = True try: if should_autolog: # This may generate warnings due to collisions in already-logged param names log_fn_args_as_params(original, args, kwargs) # training model model = original(self, *args, **kwargs) if should_autolog: # Log the model if get_autologging_config(FLAVOR_NAME, "log_models", True): global _save_model_called_from_autolog _save_model_called_from_autolog = True registered_model_name = get_autologging_config( FLAVOR_NAME, "registered_model_name", None ) try: log_model( model, artifact_path="model", registered_model_name=registered_model_name, ) finally: _save_model_called_from_autolog = False # Log the most common metrics if isinstance(model, statsmodels.base.wrapper.ResultsWrapper): metrics_dict = _get_autolog_metrics(model) mlflow.log_metrics(metrics_dict) model_summary = model.summary().as_text() mlflow.log_text(model_summary, "model_summary.txt") return model finally: # Clean the shared flag for future calls in case it had been set here ... if should_autolog: AutologHelpers.should_autolog = True
def fit(self, *args, **kwargs): should_autolog = False if AutologHelpers.should_autolog: AutologHelpers.should_autolog = False should_autolog = True if not mlflow.active_run(): try_mlflow_log(mlflow.start_run) auto_end_run = True else: auto_end_run = False try: # training model model = original_method(self, *args, **kwargs) if should_autolog: # Log the model try_mlflow_log(log_model, model, artifact_path="model") # Log the most common metrics if isinstance(model, statsmodels.base.wrapper.ResultsWrapper): metrics_dict = results_to_dict(model) try_mlflow_log(mlflow.log_metrics, metrics_dict) # This may generate warnings due to collisions in already-logged param names log_fn_args_as_params(original_method, args, kwargs) return model finally: # Clean the shared flag for future calls in case it had been set here ... if should_autolog: AutologHelpers.should_autolog = True # End current run if it had been created here ... if auto_end_run: try_mlflow_log(mlflow.end_run)
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
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
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 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) 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) 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
def train(original, *args, **kwargs): def record_eval_results(eval_results, metrics_logger): """ Create a callback function that records evaluation results. """ @exception_safe_function def callback(env): metrics_logger.record_metrics(dict(env.evaluation_result_list), env.iteration) eval_results.append(dict(env.evaluation_result_list)) return callback 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) # 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") run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: callback = record_eval_results(eval_results, metrics_logger) 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) # 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) metrics_logger.record_metrics( {"stopped_iteration": extra_step - 1}) metrics_logger.record_metrics( {"best_iteration": model.best_iteration}) metrics_logger.record_metrics( eval_results[model.best_iteration], 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: _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") # 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) 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(xgboost.DMatrix(input_example)) model_signature = infer_signature(input_example, model_output) return model_signature # Only log the model if the autolog() param log_models is set to True. if log_models: # Will only resolve `input_example` and `signature` if `log_models` is `True`. input_example, signature = resolve_input_example_and_signature( get_input_example, infer_model_signature, log_input_examples, log_model_signatures, _logger, ) try_mlflow_log( log_model, model, artifact_path="model", signature=signature, input_example=input_example, ) return model
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) try_mlflow_log(log_model, model, artifact_path='model') if auto_end_run: try_mlflow_log(mlflow.end_run) return model
def train(original, *args, **kwargs): def record_eval_results(eval_results, metrics_logger): """ Create a callback function that records evaluation results. """ # TODO: Remove `replace("SNAPSHOT", "dev")` once the following issue is addressed: # https://github.com/dmlc/xgboost/issues/6984 if Version(xgboost.__version__.replace("SNAPSHOT", "dev")) >= Version("1.3.0"): # In xgboost >= 1.3.0, user-defined callbacks should inherit # `xgboost.callback.TrainingCallback`: # https://xgboost.readthedocs.io/en/latest/python/callbacks.html#defining-your-own-callback # noqa class Callback( xgboost.callback.TrainingCallback, metaclass=ExceptionSafeAbstractClass, ): 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(): for metric_name, metric_values_on_each_iter in metric_dict.items(): key = "{}-{}".format(data_name, metric_name) # The last element in `metric_values_on_each_iter` corresponds to # the meric on the current iteration evaluation_result_dict[key] = metric_values_on_each_iter[-1] metrics_logger.record_metrics(evaluation_result_dict, epoch) eval_results.append(evaluation_result_dict) # Return `False` to indicate training should not stop return False return Callback() else: @exception_safe_function def callback(env): metrics_logger.record_metrics(dict(env.evaluation_result_list), env.iteration) eval_results.append(dict(env.evaluation_result_list)) return callback 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) # 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") run_id = mlflow.active_run().info.run_id with batch_metrics_logger(run_id) as metrics_logger: callback = record_eval_results(eval_results, metrics_logger) 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) # 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) metrics_logger.record_metrics({"stopped_iteration": extra_step - 1}) metrics_logger.record_metrics({"best_iteration": model.best_iteration}) metrics_logger.record_metrics(eval_results[model.best_iteration], 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: _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") # 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) 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(xgboost.DMatrix(input_example)) model_signature = infer_signature(input_example, model_output) return model_signature # Only log the model if the autolog() param log_models is set to True. if log_models: # Will only resolve `input_example` and `signature` if `log_models` is `True`. input_example, signature = resolve_input_example_and_signature( get_input_example, infer_model_signature, log_input_examples, log_model_signatures, _logger, ) try_mlflow_log( log_model, model, artifact_path="model", signature=signature, input_example=input_example, ) return model