| # mypy: ignore-errors |
| |
| import warnings |
| |
| import numpy as np |
| import pandas as pd # type: ignore[import-untyped] |
| from scipy.stats import gmean # type: ignore[import-untyped] |
| from sklearn.model_selection import train_test_split # type: ignore[import-untyped] |
| from sklearn.tree import DecisionTreeRegressor # type: ignore[import-untyped] |
| from train import AHTrain |
| |
| from torch._inductor.autoheuristic.autoheuristic_utils import CHOICE_COL, FEEDBACK_COL |
| |
| |
| # TODO (AlnisM): Fix these warnings |
| warnings.filterwarnings( |
| "ignore", |
| message="The behavior of DataFrame concatenation with empty or all-NA entries is deprecated", |
| ) |
| warnings.filterwarnings( |
| "ignore", |
| message="DataFrameGroupBy.apply operated on the grouping columns.", |
| ) |
| |
| |
| class AHTrainRegressionTree(AHTrain): |
| """ |
| This class is responsible for generating a heuristic by using data collected with AutoHeuristic. It will learn a |
| regression tree that predicts a score that represents how well a specific choice will perform given an input. |
| A higher score means a better choice. The heuristic will be generated in a file named <heuristic_name>.py in the |
| torch/_inductor/autoheuristic/artifacts/ directory. |
| """ |
| |
| def __init__(self): |
| super().__init__() |
| |
| 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. |
| """ |
| (df, choices, cat_feature2cats, dummy_col_2_col_val, metadata) = self.get_df( |
| log_path, nrows=nrows, apply_filters=True |
| ) |
| df_train, df_val, df_test, feature_columns = self.custom_train_test_split(df) |
| datasets = {"train": df_train, "val": df_val, "test": df_test} |
| self.add_real_datasets(datasets, other_datasets, cat_feature2cats) |
| |
| # We will do a grid search over these values |
| # Only trying out max_depths of 5, 6, and 7 because we want to keep the tree and |
| # generated code small, but smaller than 5 does not perform well enough |
| max_depths = [5, 6, 7] |
| min_samples_leafs = [1, 2, 5, 10] |
| choice_columns = [f"{CHOICE_COL}_{choice}" for choice in choices] |
| (results_df, best_model, threshold) = self.train_and_evaluate_models( |
| datasets, feature_columns, choice_columns, max_depths, min_samples_leafs |
| ) |
| |
| # prints results for all models and datasets |
| print(results_df.to_string()) |
| |
| # prints results grouped by dataset |
| for set_name in results_df["dataset"].unique(): |
| dataset_results = results_df[results_df["dataset"] == set_name] |
| dataset_results = dataset_results.sort_values(by="correct") |
| print(dataset_results.to_string() + "\n") |
| |
| feature_names = feature_columns + choice_columns |
| self.dt_to_python( |
| best_model, |
| metadata, |
| feature_names, |
| dummy_col_2_col_val, |
| heuristic_name, |
| threshold, |
| ) |
| |
| def get_df(self, log_path, cat_feature2cats=None, nrows=None, apply_filters=False): |
| """ |
| Parses the log file and processes the data into a dataframe that can be used for training. |
| """ |
| (df, metadata, feature_columns, categorical_features, choices) = self.parse_log( |
| log_path, nrows |
| ) |
| |
| def process_data( |
| df, |
| feature_columns, |
| apply_filters, |
| min_count_measurements=3, |
| max_relative_std=5, |
| ): |
| # Calculate statistics for each input and choice combination |
| def calculate_stats(group): |
| count = len(group) |
| mean = group[FEEDBACK_COL].mean() |
| std = group[FEEDBACK_COL].std() |
| relative_std = (std / mean) * 100 if mean != 0 else np.inf |
| median = group[FEEDBACK_COL].median() |
| return pd.Series( |
| { |
| "count": count, |
| "median_execution_time": median, |
| "relative_std": relative_std, |
| } |
| ) |
| |
| stats = ( |
| df.groupby(feature_columns + [CHOICE_COL]) |
| .apply(calculate_stats) |
| .reset_index() |
| ) |
| |
| if apply_filters: |
| # Remove unstables measurements |
| valid_stats = stats[ |
| (stats["count"] >= min_count_measurements) |
| & (stats["relative_std"] <= max_relative_std) |
| ] |
| # Keep only inputs with at least two valid choices |
| valid_inputs = valid_stats.groupby(feature_columns).filter( |
| lambda x: len(x) >= 2 |
| ) |
| else: |
| valid_inputs = stats |
| |
| # Compute the winner and ratios for each input |
| def get_winner_and_speedups(group): |
| mean_time = group["median_execution_time"].mean() |
| winner = group.loc[group["median_execution_time"].idxmin(), CHOICE_COL] |
| min_time = group["median_execution_time"].min() |
| max_time = group["median_execution_time"].max() |
| |
| group["winner"] = winner |
| group["speedup"] = max_time / min_time |
| group["target"] = mean_time / group["median_execution_time"] |
| |
| return group[ |
| feature_columns + [CHOICE_COL, "winner", "speedup", "target"] |
| ] |
| |
| results = ( |
| valid_inputs.groupby(feature_columns) |
| .apply(get_winner_and_speedups) |
| .reset_index(drop=True) |
| ) |
| |
| return results |
| |
| results = process_data(df, feature_columns, apply_filters) |
| (results, added_categorical_features) = self.add_new_features(results) |
| categorical_features += added_categorical_features |
| categorical_features += [CHOICE_COL] |
| |
| ( |
| 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 custom_train_test_split( |
| self, df, test_size=0.2, val_size=0.25, random_state=42 |
| ): |
| """ |
| Splits the dataframe into train, val, and test sets. |
| Also adds other datasets, specified by the user, to the train set. |
| We need to be careful, because we want to make sure that rows with the same input but different choice are |
| kept in the same set, e.g. |
| Rows that looks like this |
| input_1,choice1,... |
| input_1,choice2,... |
| should be in the same set. |
| """ |
| # We want to make sure that rows with the same input but different choice are kept in the same set |
| exclude_columns = ["speedup", "winner", "target"] |
| feature_columns = [ |
| col |
| for col in df.columns |
| if col not in exclude_columns and not col.startswith(CHOICE_COL + "_") |
| ] |
| df["input_id"] = df.groupby(feature_columns).ngroup() |
| |
| # Get unique input IDs |
| unique_inputs = df["input_id"].unique() |
| |
| # Split unique inputs into train+val and test |
| train_val_inputs, test_inputs = train_test_split( |
| unique_inputs, test_size=test_size, random_state=random_state |
| ) |
| |
| # Split train+val inputs into train and val |
| train_inputs, val_inputs = train_test_split( |
| train_val_inputs, test_size=val_size, random_state=random_state |
| ) |
| |
| # Create masks for each set |
| train_mask = df["input_id"].isin(train_inputs) |
| val_mask = df["input_id"].isin(val_inputs) |
| test_mask = df["input_id"].isin(test_inputs) |
| |
| # Split the dataframe |
| df_train = df[train_mask] |
| df_val = df[val_mask] |
| df_test = df[test_mask] |
| |
| # Remove the temporary input_id column |
| df_train = df_train.drop("input_id", axis=1) |
| df_val = df_val.drop("input_id", axis=1) |
| df_test = df_test.drop("input_id", axis=1) |
| |
| return df_train, df_val, df_test, feature_columns |
| |
| def train_and_evaluate_models( |
| self, |
| datasets, |
| feature_columns, |
| choice_columns, |
| max_depths, |
| min_samples_leafs, |
| threshold=0.99, |
| ): |
| """ |
| Does a grid search over max_depths, min_samples_leafs, and returns the best model. |
| """ |
| |
| results = [] |
| df_train = datasets["train"] |
| df_val = datasets["val"] |
| |
| best_model = None |
| best_model_threshold = 0 |
| max_correct_predictions = -1 |
| for max_depth in max_depths: |
| for min_samples_leaf in min_samples_leafs: |
| print( |
| f"Evaluating max_depth={max_depth}, min_samples_leaf={min_samples_leaf}" |
| ) |
| model = DecisionTreeRegressor( |
| random_state=42, |
| max_depth=max_depth, |
| min_samples_leaf=min_samples_leaf, |
| ) |
| model.fit( |
| df_train[feature_columns + choice_columns], df_train["target"] |
| ) |
| |
| # we first compute a safe threshold: this threshold ensures that on the validation set, |
| # if the heuristic returns a choice, the choice will be correct, although a high threshold |
| # can lead to a lot of 'unsure' choices |
| eval_result = self.evaluate_model( |
| model, df_val, feature_columns, choice_columns, threshold |
| ) |
| safe_threshold = eval_result["wrong_max_ratio"] |
| for dataset_name, dataset in datasets.items(): |
| eval_result = self.evaluate_model( |
| model, dataset, feature_columns, choice_columns, safe_threshold |
| ) |
| print(eval_result) |
| if dataset_name == "val": |
| eval_correct = eval_result["correct"] |
| if eval_correct > max_correct_predictions: |
| best_model = model |
| best_model_threshold = safe_threshold |
| max_correct_predictions = eval_correct |
| results.append( |
| { |
| "max_depth": max_depth, |
| "min_samples_leaf": min_samples_leaf, |
| "dataset": dataset_name, |
| "correct": eval_result["correct"], |
| "wrong": eval_result["wrong"], |
| "unsure": eval_result["unsure"], |
| "total": eval_result["total"], |
| "max_wrong_speedup": eval_result["max_wrong_speedup"], |
| "gman_wrong_speedup": eval_result["gman_wrong_speedup"], |
| "threshold": safe_threshold, |
| } |
| ) |
| |
| return (pd.DataFrame(results), best_model, best_model_threshold) |
| |
| def evaluate_model(self, model, df, feature_columns, choice_columns, threshold): |
| """ |
| Custom evaluation function that evaluates a learned decision tree. |
| """ |
| |
| def predict_winner(group): |
| predictions = model.predict(group[feature_columns + choice_columns]) |
| |
| # Find the index of the maximum prediction (best choice) |
| best_choice_index = np.argmax(predictions) |
| |
| # Get the corresponding choice |
| predicted_choice = ( |
| group[choice_columns].iloc[best_choice_index].idxmax().split("_")[-1] |
| ) |
| |
| # Calculate the ratio between the best and second-best prediction |
| sorted_predictions = np.sort(predictions)[::-1] |
| top_pred_ratio = ( |
| sorted_predictions[0] / sorted_predictions[1] |
| if len(sorted_predictions) > 1 |
| else np.inf |
| ) |
| |
| # If the best choice is not "significantly" better than the second best choice, |
| # the learned heuristic will return "unsure" |
| if top_pred_ratio <= threshold: |
| predicted_winner = "unsure" |
| else: |
| predicted_winner = predicted_choice |
| |
| actual_winner = group["winner"].iloc[0] |
| is_correct = ( |
| predicted_winner == actual_winner |
| if predicted_winner != "unsure" |
| else "unsure" |
| ) |
| |
| return pd.Series( |
| { |
| "predicted_winner": predicted_winner, |
| "ratio": top_pred_ratio, |
| "actual_winner": actual_winner, |
| "is_correct": is_correct, |
| "speedup": group["speedup"].iloc[ |
| 0 |
| ], # Speedup is the same for all rows in the group |
| } |
| ) |
| |
| results = df.groupby(feature_columns).apply(predict_winner).reset_index() |
| correct = (results["is_correct"].eq(True)).sum() |
| unsure = (results["is_correct"] == "unsure").sum() |
| wrong_results = results[results["is_correct"].eq(False)] |
| wrong = len(wrong_results) |
| |
| # Calculate max and geometric mean of speedup for wrong predictions |
| # Used for debugging purposes |
| wrong_speedups = wrong_results["speedup"] |
| max_wrong_speedup = wrong_speedups.max() if not wrong_speedups.empty else np.nan |
| geo_mean_wrong_speedup = ( |
| gmean(wrong_speedups) if not wrong_speedups.empty else np.nan |
| ) |
| wrong_max_ratio = wrong_results["ratio"].max() |
| |
| total = correct + wrong + unsure |
| return { |
| "correct": correct, |
| "wrong": wrong, |
| "unsure": unsure, |
| "total": total, |
| "max_wrong_speedup": max_wrong_speedup, |
| "gman_wrong_speedup": geo_mean_wrong_speedup, |
| "wrong_max_ratio": wrong_max_ratio, |
| } |
| |
| def dt_to_python( |
| self, |
| dt, |
| metadata, |
| feature_names, |
| dummy_col_2_col_val, |
| heuristic_name, |
| threshold, |
| unsafe_leaves=None, |
| ): |
| tree_ = dt.tree_ |
| feature_name = [ |
| feature_names[i] if i != -1 else "undefined!" for i in tree_.feature |
| ] |
| |
| 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, |
| dt, |
| ) |
| ) |
| fn_def = f"\n {self.gen_predict_fn_def()}" |
| lines.append(fn_def) |
| |
| def dt_to_python(node, depth): |
| indent = " " * (depth + 1) |
| false_predicate = "" |
| if tree_.feature[node] != -2: |
| name = feature_name[node] |
| threshold = tree_.threshold[node] |
| if name in dummy_col_2_col_val: |
| (orig_name, value) = dummy_col_2_col_val[name] |
| predicate = f"{indent}if str(context.get_value('{orig_name}')) != '{value}':" |
| assert ( |
| threshold == 0.5 |
| ), f"expected threshold to be 0.5 but is {threshold}" |
| else: |
| predicate = ( |
| f"{indent}if context.get_value('{name}') <= {threshold}:" |
| ) |
| lines.append(predicate) |
| dt_to_python(tree_.children_left[node], depth + 1) |
| lines.append(f"{indent}else:") |
| dt_to_python(tree_.children_right[node], depth + 1) |
| else: |
| lines.append(self.handle_leaf(tree_, node, indent, unsafe_leaves)) |
| |
| dt_to_python(0, 1) |
| |
| self.write_heuristic_to_file(lines, heuristic_name) |
| |
| def handle_leaf(self, tree_, node, indent, unsafe_leaves): |
| """ |
| Generates the code for a leaf node. This is just the value predicted by the regression tree. |
| """ |
| value = tree_.value[node][0][0] |
| return f"{indent}return {str(value)}" |
| |
| def gen_predict_fn_def(self): |
| return "def predict(self, context: AHContext) -> float:" |
| |
| 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 torch._inductor.autoheuristic.autoheuristic_utils import AHContext, AHMetadata, Choice, CHOICE_COL |
| from torch._inductor.autoheuristic.learnedheuristic_interface import ( |
| LearnedHeuristicRegression, |
| ) |
| |
| |
| class {heuristic_name}(LearnedHeuristicRegression): |
| |
| def __init__(self) -> None: |
| pass |
| |
| {self.gen_precondition(opt_name, shared_memory, device_capa)} |
| |
| def get_feedback(self, context: AHContext, choice: Choice) -> float: |
| context.context_dict[CHOICE_COL] = choice |
| return self.predict(context) |
| |
| def get_confidence_threshold(self) -> float: |
| return {threshold} |
| |
| def get_name(self) -> str: |
| return '{opt_name}'""" |
| return boiler_plate |
| |
| |
| if __name__ == "__main__": |
| train = AHTrain() |
| train.generate_heuristic() |
