From 8ad36a1763faf84bbb7405a3f335df989b9d3844 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sat, 3 Aug 2019 00:14:27 +0800 Subject: [PATCH] [AutoTVM] Fix hang/crash issues on feature extraction (#3689) * [AutoTVM] Fix hang/crash issues on feature extraction * Update xgboost_cost_model.py * fix lint --- python/tvm/autotvm/tuner/xgboost_cost_model.py | 128 +++++++++++++++---------- src/autotvm/touch_extractor.cc | 4 +- 2 files changed, 79 insertions(+), 53 deletions(-) diff --git a/python/tvm/autotvm/tuner/xgboost_cost_model.py b/python/tvm/autotvm/tuner/xgboost_cost_model.py index e278957..2653651 100644 --- a/python/tvm/autotvm/tuner/xgboost_cost_model.py +++ b/python/tvm/autotvm/tuner/xgboost_cost_model.py @@ -312,9 +312,16 @@ class XGBoostCostModel(CostModel): for i, fea in zip(need_extract, feas): fea_cache[i] = fea - ret = np.empty((len(indexes), fea_cache[indexes[0]].shape[-1]), dtype=np.float32) + feature_len = None + for idx in indexes: + if fea_cache[idx] is not None: + feature_len = fea_cache[idx].shape[-1] + break + + ret = np.empty((len(indexes), feature_len), dtype=np.float32) for i, ii in enumerate(indexes): - ret[i, :] = fea_cache[ii] + t = fea_cache[ii] + ret[i, :] = t if t is not None else 0 return ret def __del__(self): @@ -327,71 +334,88 @@ _extract_task = None def _extract_itervar_feature_index(index): """extract iteration var feature for an index in extract_space""" - config = _extract_space.get(index) - with _extract_target: - sch, args = _extract_task.instantiate(config) - fea = feature.get_itervar_feature_flatten(sch, args, take_log=True) - fea = np.concatenate((fea, list(config.get_other_option().values()))) - return fea + try: + config = _extract_space.get(index) + with _extract_target: + sch, args = _extract_task.instantiate(config) + fea = feature.get_itervar_feature_flatten(sch, args, take_log=True) + fea = np.concatenate((fea, list(config.get_other_option().values()))) + return fea + except Exception: # pylint: disable=broad-except + return None def _extract_itervar_feature_log(arg): """extract iteration var feature for log items""" - inp, res = arg - config = inp.config - with inp.target: - sch, args = inp.task.instantiate(config) - fea = feature.get_itervar_feature_flatten(sch, args, take_log=True) - x = np.concatenate((fea, list(config.get_other_option().values()))) - - if res.error_no == 0: - y = inp.task.flop / np.mean(res.costs) - else: - y = 0.0 - return x, y + try: + inp, res = arg + config = inp.config + with inp.target: + sch, args = inp.task.instantiate(config) + fea = feature.get_itervar_feature_flatten(sch, args, take_log=True) + x = np.concatenate((fea, list(config.get_other_option().values()))) + + if res.error_no == 0: + y = inp.task.flop / np.mean(res.costs) + else: + y = 0.0 + return x, y + except Exception: # pylint: disable=broad-except + return None def _extract_knob_feature_index(index): """extract knob feature for an index in extract_space""" - config = _extract_space.get(index) - return config.get_flatten_feature() + try: + config = _extract_space.get(index) + return config.get_flatten_feature() + except Exception: # pylint: disable=broad-except + return None def _extract_knob_feature_log(arg): """extract knob feature for log items""" - inp, res = arg - config = inp.config - x = config.get_flatten_feature() - - if res.error_no == 0: - with inp.target: # necessary, for calculating flops of this task - inp.task.instantiate(config) - y = inp.task.flop / np.mean(res.costs) - else: - y = 0.0 - return x, y + try: + inp, res = arg + config = inp.config + x = config.get_flatten_feature() + + if res.error_no == 0: + with inp.target: # necessary, for calculating flops of this task + inp.task.instantiate(config) + y = inp.task.flop / np.mean(res.costs) + else: + y = 0.0 + return x, y + except Exception: # pylint: disable=broad-except + return None def _extract_curve_feature_index(index): """extract sampled curve feature for an index in extract_space""" - config = _extract_space.get(index) - with _extract_target: - sch, args = _extract_task.instantiate(config) - fea = feature.get_buffer_curve_sample_flatten(sch, args, sample_n=20) - fea = np.concatenate((fea, list(config.get_other_option().values()))) - return np.array(fea) + try: + config = _extract_space.get(index) + with _extract_target: + sch, args = _extract_task.instantiate(config) + fea = feature.get_buffer_curve_sample_flatten(sch, args, sample_n=20) + fea = np.concatenate((fea, list(config.get_other_option().values()))) + return np.array(fea) + except Exception: # pylint: disable=broad-except + return None def _extract_curve_feature_log(arg): """extract sampled curve feature for log items""" - inp, res = arg - config = inp.config - with inp.target: - sch, args = inp.task.instantiate(config) - fea = feature.get_buffer_curve_sample_flatten(sch, args, sample_n=20) - x = np.concatenate((fea, list(config.get_other_option().values()))) - - if res.error_no == 0: - y = inp.task.flop / np.mean(res.costs) - else: - y = 0.0 - return x, y - + try: + inp, res = arg + config = inp.config + with inp.target: + sch, args = inp.task.instantiate(config) + fea = feature.get_buffer_curve_sample_flatten(sch, args, sample_n=20) + x = np.concatenate((fea, list(config.get_other_option().values()))) + + if res.error_no == 0: + y = inp.task.flop / np.mean(res.costs) + else: + y = 0.0 + return x, y + except Exception: # pylint: disable=broad-except + return None def custom_callback(stopping_rounds, metric, fevals, evals=(), log_file=None, maximize=False, verbose_eval=True): diff --git a/src/autotvm/touch_extractor.cc b/src/autotvm/touch_extractor.cc index 002b970..d40ebd6 100644 --- a/src/autotvm/touch_extractor.cc +++ b/src/autotvm/touch_extractor.cc @@ -131,7 +131,9 @@ void TouchExtractor::ExitItervar_() { } itervar_stack_.pop_back(); - topdown_product_ /= itervar_map[var].length; + int64_t length = itervar_map[var].length; + if (length != 0) + topdown_product_ /= length; int64_t bottomup_product = -1; for (auto kv : itervar_map[var].touch_feature) { bottomup_product = std::max(bottomup_product, kv.second.count * kv.second.reuse); -- 2.7.4