| # mypy: ignore-errors |
| |
| import itertools |
| import json |
| import logging |
| import math |
| import warnings |
| |
| |
| warnings.filterwarnings( |
| "ignore", |
| message="The behavior of DataFrame concatenation with empty or all-NA entries is deprecated", |
| ) |
| |
| from dataclasses import dataclass |
| |
| import numpy as np |
| import pandas as pd # type: ignore[import-untyped] |
| from ah_tree import DecisionTree |
| from scipy.stats import gmean |
| from sklearn.model_selection import train_test_split |
| from sklearn.tree import DecisionTreeClassifier |
| from train import AHTrain |
| |
| |
| log = logging.getLogger(__name__) |
| DEBUG = True |
| if DEBUG: |
| ch = logging.StreamHandler() |
| ch.setLevel(logging.DEBUG) |
| formatter = logging.Formatter( |
| "%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" |
| ) |
| ch.setFormatter(formatter) |
| log.addHandler(ch) |
| |
| |
| class AHTrainDecisionTree(AHTrain): |
| def __init__(self): |
| super().__init__() |
| |
| def debug_time(self, row, top_k_choices): |
| choices_feedback = json.loads(row["choice2time"]) |
| timings = sorted(choices_feedback.items(), key=lambda x: x[1]) |
| for choice, time in timings: |
| result = f"{choice} {time}" |
| if choice in top_k_choices: |
| result += " TOPK" |
| print(result) |
| |
| def is_unsafe_leaf(self, row, predicted_config, choice2time): |
| """ |
| Can be overridden by subclasses to define their own logic for deciding when a leaf is unsafe. Returns a sample |
| that landed in the leaf, the choice predicted by the tree, and a dictionary that maps each choice to the |
| execution time. One can for example decide to mark a leaf as unsafe if the predicted choice is 2x slower |
| than the fastest choice. |
| If a leaf is unsafe, the learned heuristic will always return 'unsure' if an input lands in that leaf. |
| """ |
| |
| return False |
| |
| def get_unsafe_leaves(self, model, df, feature_columns): |
| """ |
| Given a trained decision tree, and a dataframe containing the training data, returns a list of unsafe leaves. |
| """ |
| X = df[feature_columns] |
| y = df["winner"] |
| leaf_ids = model.apply(X) |
| unique_leaves = np.unique(leaf_ids) |
| |
| unsafe_leaves = [] |
| # Iterate over each leaf |
| for leaf in unique_leaves: |
| leaf_mask = leaf_ids == leaf |
| # Get samples that land in this leaf |
| leaf_X = X[leaf_mask] |
| |
| predicted_config = model.predict(leaf_X.iloc[[0]])[0] |
| |
| # For each sample, check if we should mark the leaf as unsafe |
| for idx, row in leaf_X.iterrows(): |
| choice2time = json.loads(df.loc[idx, "choice2time"]) |
| if self.is_unsafe_leaf(row, predicted_config, choice2time): |
| unsafe_leaves.append(leaf) |
| break |
| return unsafe_leaves |
| |
| def get_allowed_wrong_prediction_pct(self): |
| """ |
| This is used to determine a threshold for when a learned heuristic returns 'unsure'. |
| If this function returns 0.01, we will set the probability required for the decision tree to return a decision |
| such that at most 1% of the predictions will be wrong on the validation set. |
| """ |
| return 0.01 |
| |
| def get_grid_search_values(self): |
| """ |
| Standard values for grid search. Can be overriden. |
| """ |
| return { |
| "max_depth": [5, 6, 7], |
| "min_samples_leaf": [1, 5, 10, 0.01, 0.05, 0.02], |
| "criterion": ["gini", "entropy"], |
| } |
| |
| def predict(self, model, df, feature_columns): |
| """ |
| Returns the predictions, probabilities, and leaf ids for a given dataframe. |
| """ |
| predictions = model.predict(df[feature_columns]) |
| proba = model.predict_proba(df[feature_columns]) |
| leaf_ids = model.apply(df[feature_columns]) |
| return predictions, proba, leaf_ids |
| |
| def ranking_num_choices(self): |
| # if the heuristic is used for ranking, this function returns the number |
| # of choices that the heuristic will return |
| if self.args.ranking is None: |
| return 5 |
| return self.args.ranking |
| |
| def train_and_evaluate_models( |
| self, |
| datasets, |
| max_depths, |
| min_samples_leafs, |
| criterion_list, |
| feature_columns, |
| ranking=False, |
| ): |
| """ |
| Does a grid search over max_depths, min_samples_leafs, and criterion_list and returns the best model. |
| """ |
| |
| results = [] |
| best_model = None |
| best_model_safe_proba = 0 |
| best_model_num_correct = 0 |
| best_model_num_wrong = 0 |
| best_model_unsafe_leaves = [] |
| columns = ["set", "crit", "max_depth", "min_samples_leaf"] |
| metrics_columns = [] |
| for max_depth, min_samples_leaf, criterion in itertools.product( |
| max_depths, min_samples_leafs, criterion_list |
| ): |
| print( |
| f"max_depth={max_depth} min_samples_leaf={min_samples_leaf} criterion={criterion}" |
| ) |
| model = DecisionTreeClassifier( |
| max_depth=max_depth, |
| min_samples_leaf=min_samples_leaf, |
| criterion=criterion, |
| random_state=42, |
| ) |
| df_train = datasets["train"] |
| df_val = datasets["val"] |
| if ranking: |
| model.fit( |
| df_train[feature_columns], |
| df_train["winner"], |
| sample_weight=df_train["relative_performance"], |
| ) |
| else: |
| model.fit(df_train[feature_columns], df_train["winner"]) |
| |
| model = DecisionTree(model, feature_columns) |
| |
| if ranking: |
| model.prune(df_train, "winner", k=self.ranking_num_choices()) |
| |
| unsafe_leaves = self.get_unsafe_leaves(model, df_train, feature_columns) |
| predictions, proba, leaf_ids = self.predict(model, df_val, feature_columns) |
| |
| wrong_pct = self.get_allowed_wrong_prediction_pct() |
| evaluator = DecisionEvaluator( |
| self, |
| model, |
| predictions, |
| df_val, |
| proba, |
| wrong_pct=wrong_pct, |
| unsafe_leaves=unsafe_leaves, |
| leaf_ids=leaf_ids, |
| k=self.ranking_num_choices(), |
| ranking=ranking, |
| ) |
| safe_proba = evaluator.get_safe_proba() |
| print(f"safe_proba={safe_proba}") |
| |
| def eval(name, df): |
| if ranking: |
| # when ranking is enabled, we duplicate each input for each choice that |
| # is almost as good as the best choice |
| # we do not want to evaluate the same input multiple times, so we remove duplicates here |
| df = df[df["winner"] == df["actual_winner"]] |
| predictions, proba, leaf_ids = self.predict(model, df, feature_columns) |
| evaluator = DecisionEvaluator( |
| self, |
| model, |
| predictions, |
| df, |
| proba, |
| wrong_pct=wrong_pct, |
| threshold=safe_proba, |
| unsafe_leaves=unsafe_leaves, |
| leaf_ids=leaf_ids, |
| k=self.ranking_num_choices(), |
| ranking=ranking, |
| ) |
| return evaluator.get_results() |
| |
| for dataset_name, dataset in datasets.items(): |
| eval_result: EvalResults = eval(dataset_name, dataset) |
| eval_result_metrics = eval_result.to_map() |
| if dataset_name == "val": |
| num_correct = eval_result.accuracy.num_correct |
| num_wrong = eval_result.accuracy.num_wrong |
| num_total = eval_result.accuracy.total |
| if num_wrong <= num_total * wrong_pct: |
| if num_correct > best_model_num_correct: |
| print( |
| f"new best model with {num_correct} correct and {num_wrong} wrong" |
| ) |
| best_model = model |
| best_model_num_correct = num_correct |
| best_model_num_wrong = num_wrong |
| best_model_safe_proba = safe_proba |
| best_model_unsafe_leaves = unsafe_leaves |
| |
| result = (dataset_name, criterion, max_depth, min_samples_leaf) |
| result += tuple(eval_result_metrics.values()) |
| results.append(result) |
| if len(metrics_columns) == 0: |
| metrics_columns = list(eval_result_metrics.keys()) |
| columns += metrics_columns |
| |
| return ( |
| pd.DataFrame(results, columns=columns), |
| best_model, |
| best_model_safe_proba, |
| best_model_unsafe_leaves, |
| ) |
| |
| def get_test_and_val_size(self): |
| """ |
| Returns the size of the test and validation sets. |
| """ |
| return (0.15, 0.15) |
| |
| def prepare_datasets(self, df, other_datasets, cat_feature2cats, ranking=False): |
| """ |
| Splits the dataframe into train, val, and test sets. |
| Also adds other datasets, specified by the user, to the train set. |
| """ |
| test_size, val_size = self.get_test_and_val_size() |
| # Split into train+val and test |
| df_train_val, df_test = train_test_split( |
| df, test_size=test_size, random_state=42 |
| ) |
| |
| # Split train+val inputs into train and val |
| train_val_size = 1 - test_size |
| df_train, df_val = train_test_split( |
| df_train_val, test_size=val_size / train_val_size, random_state=42 |
| ) |
| datasets = {"train": df_train, "val": df_val, "test": df_test} |
| self.add_real_datasets(datasets, other_datasets, cat_feature2cats, ranking) |
| return datasets |
| |
| def export_to_dot(self, best_model, df, feature_columns): |
| """ |
| Export a learned decision tree to a dot file. |
| """ |
| dot_str = best_model.to_dot() |
| with open("best_model.dot", "w") as f: |
| f.write(dot_str) |
| |
| def get_feature_columns(self, df): |
| """ |
| The dataframe contains columns that are not features, such as 'winner', 'speedup' that are only used for |
| debugging purposes. This function returns the columns that are actually features. |
| """ |
| exclude_columns = [ |
| "speedup", |
| "winner", |
| "target", |
| "avail_choices", |
| "choice2time", |
| "index", |
| "actual_winner", |
| "relative_performance", |
| ] |
| feature_columns = [col for col in df.columns if col not in exclude_columns] |
| return feature_columns |
| |
| def add_training_data(self, df_train, datasets): |
| return datasets["train"] |
| |
| def main( |
| self, |
| log_path, |
| other_datasets, |
| nrows, |
| heuristic_name, |
| save_dot=False, |
| ranking=False, |
| ): |
| """ |
| Main function that trains a decision tree and generates a heuristic. |
| """ |
| # TODO: Enable apply_filters |
| (df, choices, cat_feature2cats, dummy_col_2_col_val, metadata) = self.get_df( |
| log_path, nrows=nrows, apply_filters=False, add_near_best=ranking |
| ) |
| self.dummy_col_2_col_val = dummy_col_2_col_val |
| datasets = self.prepare_datasets(df, other_datasets, cat_feature2cats, ranking) |
| df_train = self.add_training_data(datasets["train"], datasets) |
| datasets["train"] = df_train |
| print(datasets["train"]["winner"].value_counts().to_string()) |
| |
| feature_columns = self.get_feature_columns(df) |
| grid_search_values = self.get_grid_search_values() |
| max_depths = grid_search_values["max_depth"] |
| min_samples_leafs = grid_search_values["min_samples_leaf"] |
| criterion_list = grid_search_values["criterion"] |
| ( |
| results_df, |
| best_model, |
| best_model_safe_proba, |
| unsafe_leaves, |
| ) = self.train_and_evaluate_models( |
| datasets, |
| max_depths, |
| min_samples_leafs, |
| criterion_list, |
| feature_columns, |
| ranking=ranking, |
| ) |
| |
| if ranking: |
| columns_to_keep = [ |
| "set", |
| "crit", |
| "max_depth", |
| "min_samples_leaf", |
| "total", |
| "top_k_correct", |
| "top_k_wrong", |
| "top_k_unsure", |
| "wrong_max_speedup_k", |
| "wrong_gmean_speedup_k", |
| ] |
| results_df = results_df[columns_to_keep] |
| # prints results for all models and datasets |
| print(results_df.to_string()) |
| |
| sort_metric = "top_k_correct" if ranking else "correct" |
| # prints results grouped by dataset |
| for set_name in results_df["set"].unique(): |
| dataset_results = results_df[results_df["set"] == set_name] |
| dataset_results = dataset_results.sort_values(by=sort_metric) |
| print(dataset_results.to_string() + "\n") |
| |
| if best_model is not None: |
| if save_dot: |
| self.export_to_dot(best_model, df, feature_columns) |
| self.codegen( |
| best_model, |
| metadata, |
| heuristic_name, |
| best_model_safe_proba, |
| dummy_col_2_col_val, |
| unsafe_leaves, |
| ) |
| else: |
| print( |
| "All learned models have too many wrong predictions, so no heuristic was generated" |
| ) |
| |
| def get_df( |
| self, |
| log_path, |
| cat_feature2cats=None, |
| nrows=None, |
| apply_filters=False, |
| add_near_best=False, |
| ): |
| """ |
| Parses the log file and processes the data into a dataframe that can be used for training. |
| """ |
| (df, metadata, features, categorical_features, choices) = self.parse_log( |
| log_path, nrows |
| ) |
| |
| def calculate_stats(group): |
| count = len(group) |
| has_inf = np.isinf(group["feedback"]).any() |
| if has_inf: |
| relative_std = np.inf |
| median = np.inf |
| else: |
| mean = group["feedback"].mean() |
| std = group["feedback"].std() |
| relative_std = (std / mean) * 100 if mean != 0 else np.inf |
| median = group["feedback"].median() |
| if relative_std > 5: |
| times = group["feedback"].tolist() |
| times_str = ", ".join([f"{t:.3f}" for t in sorted(times)]) |
| log.debug("High relative std: %f. times=%s", relative_std, times_str) |
| return pd.Series( |
| { |
| "count": count, |
| "relative_std": relative_std, |
| "median_execution_time": median, |
| } |
| ) |
| |
| feature_columns = features |
| stats = ( |
| df.groupby(feature_columns + ["choice"], as_index=False) |
| .apply(calculate_stats, include_groups=False) |
| .reset_index() |
| ) |
| |
| # TODO: We have to be careful with removing certain choices, because if we e.g. remove the winner, the |
| # heuristic will end up learning wrong things. But, execution times with high variance are also bad |
| if apply_filters: |
| # Filter out inputs with less than 3 measurements or high relative std |
| valid_stats = stats[(stats["count"] >= 3) & (stats["relative_std"] <= 5)] |
| # Group by input features and count how many valid choices we have for each input |
| valid_inputs = valid_stats.groupby(feature_columns).filter( |
| lambda x: len(x) >= 2 |
| ) |
| else: |
| valid_inputs = stats |
| |
| # Compute the winner and speedup for each valid input |
| def get_winner_and_speedup(group): |
| assert len(group) >= 2, "Need at least 2 choices" |
| |
| sorted_group = group.sort_values("median_execution_time") |
| winner = sorted_group.iloc[0]["choice"] |
| winning_time = sorted_group.iloc[0]["median_execution_time"] |
| second_best_time = sorted_group.iloc[1]["median_execution_time"] |
| speedup = second_best_time / winning_time |
| unique_choices = group["choice"].unique() |
| |
| choice2time = {} |
| for row in group.itertuples(): |
| choice2time[row.choice] = row.median_execution_time |
| |
| assert len(unique_choices) == len( |
| group |
| ), f"len(unique_choices) != len(group): {len(unique_choices)} != {len(group)}" |
| |
| return pd.Series( |
| { |
| "winner": winner, |
| "speedup": speedup, |
| "avail_choices": unique_choices, |
| "choice2time": json.dumps(choice2time), |
| } |
| ) |
| |
| results = ( |
| valid_inputs.groupby(feature_columns, as_index=False) |
| .filter(lambda x: len(x) >= 2) |
| .groupby(feature_columns, as_index=False) |
| .apply(get_winner_and_speedup, include_groups=False) |
| .reset_index() |
| ) |
| |
| def add_near_best_configs(df): |
| new_rows = [] |
| |
| for index, row in df.iterrows(): |
| dictionary = json.loads(row["choice2time"]) |
| min_value = min(dictionary.values()) |
| |
| for key, value in dictionary.items(): |
| new_row = row.copy() |
| relative_performance = min_value / value |
| new_row["relative_performance"] = relative_performance |
| if relative_performance is None or relative_performance is np.inf: |
| breakpoint() |
| new_row["actual_winner"] = row["winner"] |
| new_row["winner"] = key |
| if relative_performance >= 0.98: |
| new_rows.append(new_row) |
| |
| return pd.DataFrame(new_rows).reset_index(drop=True) |
| |
| if add_near_best: |
| results = add_near_best_configs(results) |
| (results, added_categorical_features) = self.add_new_features(results) |
| categorical_features += added_categorical_features |
| |
| ( |
| results, |
| cat_feature2cats, |
| dummy_col_2_col_val, |
| ) = self.handle_categorical_features( |
| cat_feature2cats, categorical_features, results |
| ) |
| return (results, choices, cat_feature2cats, dummy_col_2_col_val, metadata) |
| |
| def ranking_always_included_choices(self): |
| return [] |
| |
| def gen_classes(self, classes, num_spaces): |
| """ |
| If classes=['choice1', 'choice2', 'choice3'], then this function returns |
| the following string: |
| self.choices.append('choice1') |
| self.choices.append('choice2') |
| self.choices.append('choice3') |
| Used in the generated heuristic to map the index of a choice to its name. |
| """ |
| indent = " " * num_spaces |
| return "\n".join([f"{indent}self.choices.append('{c}')" for c in classes]) |
| |
| def get_default_config(self, row): |
| """ |
| Returns the default config for a given sample. The default config could for example be the config that is |
| the chosen by a current handwritten heuristic. This can for example be used in get_unsafe_leaf to |
| compare the predicted config with the default config. |
| """ |
| return None |
| |
| def gen_predict_fn_def(self): |
| """ |
| Generates the definition of the predict function. |
| """ |
| return "def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:" |
| |
| def codegen_boilerplate( |
| self, heuristic_name, opt_name, threshold, shared_memory, device_capa, classes |
| ): |
| """ |
| Generates the boilerplate code for the generated heuristic. This includes things like imports, class definition, |
| etc. |
| """ |
| |
| boiler_plate = f"""# flake8: noqa: B950 |
| # fmt: off |
| # This file was generated by AutoHeuristic. Do not modify it manually! |
| # To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/{opt_name}/ |
| from typing import List, Optional, Tuple |
| |
| from torch._inductor.autoheuristic.autoheuristic_utils import ( |
| AHContext, |
| AHMetadata, |
| Choice, |
| ) |
| from torch._inductor.autoheuristic.learnedheuristic_interface import ( |
| LearnedHeuristicDecision, |
| ) |
| |
| |
| class {heuristic_name}(LearnedHeuristicDecision): |
| |
| def __init__(self) -> None: |
| self.choices: List[Choice] = [] |
| self.fill_choices() |
| |
| {self.gen_precondition(opt_name, shared_memory, device_capa)} |
| |
| def get_confidence_threshold(self) -> float: |
| return {threshold} |
| |
| def get_choice(self, idx: int) -> Optional[str]: |
| if idx < len(self.choices): |
| return self.choices[idx] |
| return None |
| |
| def fill_choices(self) -> None: |
| {self.gen_classes(classes, num_spaces=8)} |
| |
| def get_name(self) -> str: |
| return '{opt_name}'""" |
| return boiler_plate |
| |
| def add_real_datasets( |
| self, datasets, other_datasets, cat_feature2cats, ranking=False |
| ): |
| """ |
| Adds datasets specified by the user to the datasets dictionary. |
| """ |
| if other_datasets: |
| for name, path in other_datasets: |
| (df_other, choices, _, _, _) = self.get_df( |
| path, |
| cat_feature2cats=cat_feature2cats, |
| apply_filters=False, |
| add_near_best=ranking, |
| ) |
| datasets[name] = df_other |
| |
| def codegen( |
| self, |
| tree, |
| metadata, |
| heuristic_name, |
| threshold, |
| dummy_col_2_col_val, |
| unsafe_leaves, |
| ): |
| lines = [] |
| device_capa = metadata["device_capa"] |
| device_capa_str = f"({device_capa[0]}, {device_capa[1]})" |
| opt_name = metadata["name"] |
| lines.append( |
| self.codegen_boilerplate( |
| heuristic_name, |
| opt_name, |
| threshold, |
| metadata["shared_memory"], |
| device_capa_str, |
| tree.classes_, |
| ) |
| ) |
| fn_def = f"\n {self.gen_predict_fn_def()}" |
| lines.append(fn_def) |
| tree.codegen(dummy_col_2_col_val, lines, unsafe_leaves) |
| self.write_heuristic_to_file(lines, heuristic_name) |
| |
| |
| @dataclass |
| class AccuracyMetrics: |
| # Number of correct predictions |
| num_correct: int |
| # Number of wrong predictions |
| num_wrong: int |
| # Number of predictions where model is unsure |
| num_unsure: int |
| # Total number of predictions |
| total: int |
| |
| def to_map(self): |
| return { |
| "correct": self.num_correct, |
| "wrong": self.num_wrong, |
| "unsure": self.num_unsure, |
| "total": self.total, |
| } |
| |
| |
| @dataclass |
| class WrongSpeedupMetrics: |
| # If the model predicted the wrong choice, this is the maximum speedup of the best choice over the predicted choice |
| max_speedup: float |
| # For all wrong predictions, this is the geometric mean of the speedups of the best choices over the predicted choices |
| gmean_speedup: float |
| |
| def to_map(self): |
| return { |
| "wrong_max_speedup": self.max_speedup, |
| "wrong_gmean_speedup": self.gmean_speedup, |
| } |
| |
| |
| @dataclass |
| class RankingMetrics: |
| # Number of predictions where best choice is in top k choices |
| num_correct: int |
| # Number of predictions where best choice is not in top k choices |
| num_wrong: int |
| # Maximum speedup of best choice over best choice in top k (this tells us how much better the best choice, which |
| # is not in top k, is over the best choice in top k) |
| max_speedup: float |
| # Geometric mean of speedups of best choice over best choice in top k |
| gmean_speedup: float |
| # Number of predictions where model is unsure |
| unsure: int |
| |
| def to_map(self): |
| return { |
| "top_k_correct": self.num_correct, |
| "top_k_wrong": self.num_wrong, |
| "wrong_max_speedup_k": self.max_speedup, |
| "wrong_gmean_speedup_k": self.gmean_speedup, |
| "top_k_unsure": self.unsure, |
| } |
| |
| |
| @dataclass |
| class DefaultComparisonMetrics: |
| # Maximum speedup of predicted choice over default choice |
| max_speedup: float |
| # Geometric mean of speedups of predicted choices over default choices |
| gmean_speedup: float |
| # Maximum speedup of default choice over predicted choice |
| max_slowdown: float |
| # Number of predictions where the predicted choice is not the default choice |
| non_default_predictions: int |
| # Number of predictions where the default choice is better than the predicted choice |
| default_better: bool |
| |
| def to_map(self): |
| return { |
| "max_speedup_over_default": self.max_speedup, |
| "gmean_speedup_over_default": self.gmean_speedup, |
| "max_speedup_default_over_heuristic": self.max_slowdown, |
| "non_default_predictions": self.non_default_predictions, |
| "default_better": self.default_better, |
| } |
| |
| |
| @dataclass |
| class EvalResults: |
| accuracy: AccuracyMetrics |
| speedup: WrongSpeedupMetrics |
| ranking: RankingMetrics |
| default_comparison: DefaultComparisonMetrics |
| |
| def to_map(self): |
| return { |
| **self.accuracy.to_map(), |
| **self.speedup.to_map(), |
| **self.ranking.to_map(), |
| **self.default_comparison.to_map(), |
| } |
| |
| |
| class DecisionEvaluator: |
| def __init__( |
| self, |
| train, |
| model, |
| predictions, |
| df, |
| probas, |
| wrong_pct=0.01, |
| threshold=0.0, |
| k=10, |
| unsafe_leaves=None, |
| leaf_ids=None, |
| ranking=False, |
| ) -> None: |
| self.train = train |
| self.model = model |
| self.predictions = predictions |
| self.df = df |
| self.probas = probas |
| self.wrong_pct = wrong_pct |
| self.threshold = threshold |
| self.k = k |
| self.unsafe_leaves = unsafe_leaves |
| self.leaf_ids = leaf_ids |
| self.ranking = ranking |
| |
| self.num_correct = 0 |
| self.num_wrong = 0 |
| self.num_unsure = 0 |
| self.wrong_probas = [] |
| self.speedups_wrong = [] |
| self.num_correct_top_k = 0 |
| self.num_wrong_top_k = 0 |
| self.wrong_speedups_top_k = [] |
| self.top_k_unsure = 0 |
| self.num_non_default_predictions = 0 |
| self.speedups_over_default = [] |
| self.num_default_better = 0 |
| |
| def compute_speedup_over_default(self, default_config, pred, i, predicted_time): |
| if default_config is not None: |
| if pred != default_config: |
| self.num_non_default_predictions += 1 |
| default_time = self.get_time(self.df.iloc[i], default_config) |
| # TODO: We should keep track of how often this happens |
| if default_time is not None and not math.isinf(default_time): |
| speedup_over_default = default_time / predicted_time |
| if speedup_over_default < 1: |
| self.num_default_better += 1 |
| self.speedups_over_default.append(speedup_over_default) |
| else: |
| log.debug( |
| "cannot compute speedup over default because default_time=%d", |
| default_time, |
| ) |
| |
| def get_time(self, row, choice): |
| choices_feedback = json.loads(row["choice2time"]) |
| return choices_feedback.get(choice, None) |
| |
| def top_k_classes(self, model, probas, k, avail_choices): |
| # Get classes and their corresponding probabilities |
| classes = model.classes_ |
| class_proba_pairs = list(zip(classes, probas)) |
| |
| # Sort by probability (descending) and filter out zero probabilities |
| sorted_classes = [ |
| c |
| for c, p in sorted(zip(classes, probas), key=lambda x: x[1], reverse=True) |
| if p > 0 and c in avail_choices |
| ] |
| |
| # Return top k choices |
| top_k_choices = sorted_classes[:k] |
| top_k_choices += self.train.ranking_always_included_choices() |
| top_k_choices = list(dict.fromkeys(top_k_choices)) |
| return top_k_choices |
| |
| def eval_prediction( |
| self, avail_choices, leaf_id, pred, true, prob, threshold, default_config, i |
| ): |
| predicted_time = self.get_time(self.df.iloc[i], pred) |
| max_prob = max(prob) |
| if ( |
| leaf_id in self.unsafe_leaves |
| or pred not in avail_choices |
| or (max_prob != 1.0 and max_prob <= threshold) |
| ): |
| self.num_unsure += 1 |
| self.speedups_over_default.append(1.0) |
| elif pred == true: |
| self.compute_speedup_over_default(default_config, pred, i, predicted_time) |
| self.num_correct += 1 |
| else: |
| self.compute_speedup_over_default(default_config, pred, i, predicted_time) |
| self.num_wrong += 1 |
| self.wrong_probas.append(max_prob) |
| best_time = self.get_time(self.df.iloc[i], true) |
| wrong_speedup = predicted_time / best_time |
| self.speedups_wrong.append(wrong_speedup) |
| |
| def eval_ranking_prediction(self, true, top_k_choices, i): |
| if true in top_k_choices: |
| self.num_correct_top_k += 1 |
| else: |
| top_k_choices_times = [] |
| for choice in top_k_choices: |
| time = self.get_time(self.df.iloc[i], choice) |
| if time is not None: |
| top_k_choices_times.append(time) |
| best_time = self.get_time(self.df.iloc[i], true) |
| min_time = min(top_k_choices_times, default=None) |
| if min_time is not None: |
| speedup = min_time / best_time |
| self.wrong_speedups_top_k.append(speedup) |
| self.num_wrong_top_k += 1 |
| else: |
| self.top_k_unsure += 1 |
| # TODO (AlnisM): print more info (input and choices) |
| log.debug( |
| "All top k choices have no time which means all top k are unavailable" |
| ) |
| |
| def get_safe_proba(self): |
| return self.get_results(return_safe_proba=True) |
| |
| def compute_safe_proba(self, num_predictions, wrong_probas, wrong_pct): |
| wrong_probas.sort() |
| num_wrong = len(wrong_probas) |
| allowed_wrong = int(num_predictions * wrong_pct) |
| if allowed_wrong >= num_wrong: |
| return 0.0 |
| too_many_wrong = num_wrong - allowed_wrong |
| idx = min(too_many_wrong, len(wrong_probas) - 1) |
| return wrong_probas[idx] |
| |
| def get_results(self, return_safe_proba=False) -> EvalResults: |
| """ |
| Custom evaluation function that evaluates a learned decision tree. |
| """ |
| |
| y_true = self.df["actual_winner"] if self.ranking else self.df["winner"] |
| i = 0 |
| for pred, true, prob, leaf_id in zip( |
| self.predictions, y_true, self.probas, self.leaf_ids |
| ): |
| avail_choices = self.df["avail_choices"].iloc[i] |
| top_k_choices = self.top_k_classes( |
| self.model, prob, k=self.k, avail_choices=avail_choices |
| ) |
| assert ( |
| true in avail_choices |
| ), f"Best choice {true} not in available choices {avail_choices}" |
| default_config = self.train.get_default_config(self.df.iloc[i]) |
| self.eval_prediction( |
| avail_choices, |
| leaf_id, |
| pred, |
| true, |
| prob, |
| self.threshold, |
| default_config, |
| i, |
| ) |
| self.eval_ranking_prediction(true, top_k_choices, i) |
| i += 1 |
| |
| total = len(self.predictions) |
| if return_safe_proba: |
| return self.compute_safe_proba(total, self.wrong_probas, self.wrong_pct) |
| |
| def safe_gmean(x): |
| return gmean(x) if x else 0 |
| |
| max_speedup = max(self.speedups_wrong, default=0) |
| gmean_speedup = safe_gmean(self.speedups_wrong) |
| max_speedup_top_k = max(self.wrong_speedups_top_k, default=0) |
| gmean_speedup_top_k = safe_gmean(self.wrong_speedups_top_k) |
| max_speedup_over_default = max(self.speedups_over_default, default=0) |
| gmean_speedup_over_default = safe_gmean(self.speedups_over_default) |
| max_slowdown_over_default = min(self.speedups_over_default, default=0) |
| |
| accuracyMetrics = AccuracyMetrics( |
| self.num_correct, self.num_wrong, self.num_unsure, total |
| ) |
| wrongSpeedupMetrics = WrongSpeedupMetrics(max_speedup, gmean_speedup) |
| rankingMetrics = RankingMetrics( |
| self.num_correct_top_k, |
| self.num_wrong_top_k, |
| max_speedup_top_k, |
| gmean_speedup_top_k, |
| self.top_k_unsure, |
| ) |
| defaultComparisonMetrics = DefaultComparisonMetrics( |
| max_speedup_over_default, |
| gmean_speedup_over_default, |
| max_slowdown_over_default, |
| self.num_non_default_predictions, |
| self.num_default_better, |
| ) |
| return EvalResults( |
| accuracyMetrics, |
| wrongSpeedupMetrics, |
| rankingMetrics, |
| defaultComparisonMetrics, |
| ) |
| |
| |
| if __name__ == "__main__": |
| train = AHTrainDecisionTree() |
| train.generate_heuristic() |
