| """ | |
| Copyright (c) 2024, Alliance for Open Media. All rights reserved | |
| This source code is subject to the terms of the BSD 3-Clause Clear License | |
| and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear | |
| License was not distributed with this source code in the LICENSE file, you | |
| can obtain it at aomedia.org/license/software-license/bsd-3-c-c/. If the | |
| Alliance for Open Media Patent License 1.0 was not distributed with this | |
| source code in the PATENTS file, you can obtain it at | |
| aomedia.org/license/patent-license/. | |
| """ | |
| import json | |
| import os | |
| import numpy as np | |
| import pandas as pd | |
| import parakit.entropy.model as model | |
| from parakit.entropy.codec_cdf_functions import ( | |
| cdf2pmf_av1, | |
| cdfinv2pmf_av1, | |
| cost_symbol_av1, | |
| pmf2cdfinv_av1, | |
| update_cdfinv_av1, | |
| ) | |
| from parakit.entropy.data_collector import DataCollector | |
| COST_REGULARIZATION = 4 # modifies maximum cost | |
| MIN_SAMPLE_REQUIRED = 10 # minimum number of data points needed to run training | |
| MAX_NUMBER_ROWS = 5000000 | |
| class Trainer: | |
| def __init__( | |
| self, | |
| filename_data, | |
| rate_search_list, | |
| cost_regularization=COST_REGULARIZATION, | |
| min_sample=MIN_SAMPLE_REQUIRED, | |
| max_rows=MAX_NUMBER_ROWS, | |
| ): | |
| # training parameters | |
| self._filename_data = filename_data # data file | |
| self._cost_regularization = cost_regularization | |
| self._min_sample = min_sample | |
| self._max_rows = max_rows | |
| self._rate_search_list = rate_search_list | |
| # internal data and models | |
| self.dataframe = None | |
| self.initial_model = None | |
| self.resulting_model = None | |
| self._max_cost_regularization = None | |
| self._output_filename = None # resulting json file | |
| self.initialize() | |
| def initialize(self): | |
| dc = DataCollector(self._filename_data) | |
| self._initial_model = dc.get_context_model() | |
| self.dataframe = dc.collect_dataframe(max_rows=self._max_rows) | |
| self._output_filename = self._get_output_filename() | |
| self._max_cost_regularization = cost_symbol_av1(self._cost_regularization) | |
| def prepare_dataframe(self, df, ctx_idx_list): | |
| dims = self._initial_model.num_dims | |
| df_filtered = df | |
| if dims > 0: | |
| mask = df["Dim0"] == ctx_idx_list[0] | |
| for i in range(1, dims): | |
| mask &= df[f"Dim{i}"] == ctx_idx_list[i] | |
| df_filtered = df.loc[mask] | |
| # add "CalcCost" column calculating cost | |
| new_cols = ["CalcCost"] | |
| df_filtered = df_filtered.reindex( | |
| columns=[*df_filtered.columns.tolist(), *new_cols], fill_value=0 | |
| ) | |
| df_filtered.reset_index(inplace=True) | |
| # return empty df if very few (< _min_sample) data points available | |
| if len(df_filtered) < self._min_sample: | |
| return pd.DataFrame() | |
| # convert NaN values to -1, if any | |
| df_filtered = df_filtered.fillna(-1) | |
| df_filtered = df_filtered[df_filtered.columns].astype( | |
| int | |
| ) # all columns are integer | |
| return df_filtered | |
| def map_rate_from_counter(self, rates_tuple, counter): | |
| if counter <= 15: | |
| return rates_tuple[0] | |
| elif counter <= 31: | |
| return rates_tuple[1] | |
| else: | |
| return rates_tuple[2] | |
| def update_pmf_per_rate(self, pmf_list, cost_list, val, counter): | |
| rate_offset_list = self._rate_search_list | |
| nsymb = self._initial_model.num_symb | |
| for r, rates_tuple in enumerate(rate_offset_list): | |
| roffset = self.map_rate_from_counter(rates_tuple, counter) | |
| # cost | |
| cost = cost_symbol_av1(pmf_list[r][val]) | |
| if cost > self._max_cost_regularization: | |
| cost = self._max_cost_regularization | |
| cost_list[r] += cost | |
| # update | |
| cdf_inv = pmf2cdfinv_av1(pmf_list[r]) | |
| cdf_inv_updated = update_cdfinv_av1(cdf_inv, val, counter, nsymb, roffset) | |
| pmf_list[r] = cdfinv2pmf_av1(cdf_inv_updated) | |
| return (pmf_list, cost_list) | |
| def estimate_cost_multiple(self, df_filtered, pmf_init, rateidx_init=0): | |
| nsymb = self._initial_model.num_symb | |
| rate_offset_list = self._rate_search_list | |
| num_rates = len(rate_offset_list) | |
| cost_list = np.zeros(num_rates, dtype=int) | |
| # initialization | |
| pmf_init_list = [pmf_init] * num_rates | |
| prev_pmf_buffer = [pmf_init_list] | |
| cdf_cols = [] | |
| for i in range(nsymb): | |
| if f"cdf{i}" in df_filtered.keys(): | |
| cdf_cols.append(f"cdf{i}") | |
| mask_index = (df_filtered["Counter"] == 0) | (df_filtered["isBeginFrame"] == 1) | |
| index_list = df_filtered[mask_index].index.to_list() | |
| index_list.append(len(df_filtered)) | |
| index_pair_list = [ | |
| (index_list[ind], index_list[ind + 1]) for ind in range(len(index_list) - 1) | |
| ] | |
| for index_pair in index_pair_list: | |
| start_idx, end_idx = index_pair | |
| df_sub = df_filtered.iloc[start_idx:end_idx].copy() | |
| num_data_samples = len(df_sub) | |
| success = False | |
| if len(cdf_cols) == nsymb: | |
| cdf_init = np.array(df_sub.iloc[0][cdf_cols], dtype=int) | |
| prev_pmf_buffer = [[cdf2pmf_av1(cdf_init)] * num_rates] | |
| for _, prev_pmf in enumerate(prev_pmf_buffer): | |
| rateidx_init = df_sub.iloc[0]["rate"] | |
| pmf_default = prev_pmf[rateidx_init] # default one | |
| # initial values | |
| val_init = df_sub.iloc[0]["Value"] | |
| cost = df_sub.iloc[0]["Cost"] | |
| calculated_cost = cost_symbol_av1(pmf_default[val_init]) | |
| if calculated_cost > self._max_cost_regularization: | |
| calculated_cost = self._max_cost_regularization | |
| cost_list_persub = np.zeros(num_rates, dtype=int) | |
| if calculated_cost != cost: | |
| continue | |
| pmf_list = [pmf_default] * num_rates | |
| # try best cases | |
| for i in range(num_data_samples): | |
| val = df_sub.iloc[i]["Value"] | |
| counter = df_sub.iloc[i]["Counter"] | |
| cost = df_sub.iloc[i]["Cost"] | |
| calculated_cost = cost_symbol_av1(pmf_list[rateidx_init][val]) | |
| if calculated_cost > self._max_cost_regularization: | |
| calculated_cost = self._max_cost_regularization | |
| df_sub.iloc[i]["CalcCost"] = calculated_cost | |
| if cost != calculated_cost: | |
| break | |
| # update variables | |
| pmf_list, cost_list_persub = self.update_pmf_per_rate( | |
| pmf_list, cost_list_persub, val, counter | |
| ) | |
| # success | |
| if i == num_data_samples - 1: | |
| success = True | |
| # update cost | |
| cost_list += cost_list_persub | |
| # save latest pmf before updating | |
| prev_pmf_buffer.insert(1, pmf_list) | |
| if success: | |
| break | |
| return (cost_list, 0) | |
| def get_cost_per_value(self, df): | |
| num_symb = self._initial_model.num_symb | |
| cost_pervalue_list = [0] * num_symb | |
| if df.empty: | |
| return cost_pervalue_list | |
| for i in range(num_symb): | |
| cost = df[df["Value"] == i]["Cost"].sum() | |
| cost_pervalue_list[i] = int(cost) | |
| return cost_pervalue_list | |
| def get_value_count(self, df): | |
| num_symb = self._initial_model.num_symb | |
| value_count_list = [0] * num_symb | |
| if df.empty: | |
| return value_count_list | |
| count_series = df["Value"].value_counts() | |
| for i in count_series.index: | |
| count = count_series[i] | |
| value_count_list[i] = int(count) | |
| return value_count_list | |
| def run_rate_training_on_file(self): | |
| # parameters | |
| rate_parameters = self._rate_search_list | |
| size_list = self._initial_model.size_list | |
| num_dims = self._initial_model.num_dims | |
| num_symb = self._initial_model.num_symb | |
| ctx_group_name = self._initial_model.ctx_group_name | |
| prob_dict = self._initial_model.model_dict | |
| ctx_name_list = list(prob_dict.keys()) | |
| result_dict = prob_dict.copy() | |
| for ctx_name in ctx_name_list: | |
| ctx_idx_interest = prob_dict[ctx_name]["index"] | |
| print(f"Context {ctx_name}:", end=" ") | |
| df_filtered = self.prepare_dataframe(self.dataframe, ctx_idx_interest) | |
| value_count_list = self.get_value_count(df_filtered) | |
| value_cost_list = self.get_cost_per_value(df_filtered) | |
| # check empty df | |
| if df_filtered.empty: | |
| print("Skip training - insufficient data.") | |
| result_dict[ctx_name] = { | |
| "current_cost": np.zeros(len(rate_parameters), dtype=int), | |
| "adapt_rate": rate_parameters, | |
| "initial_cost": 0, | |
| "codec_cost": 0, | |
| "upper_cost": 0, | |
| "num_samples": 0, | |
| "initializer": prob_dict[ctx_name]["initializer"], | |
| "init_rateidx": prob_dict[ctx_name]["init_rateidx"], | |
| "value_count": value_count_list, | |
| "value_cost": value_cost_list, | |
| } | |
| continue | |
| # actual cost | |
| codec_cost = df_filtered["Cost"].sum() | |
| # estimate cost | |
| default_cdf = np.array(prob_dict[ctx_name]["initializer"]) | |
| default_pmf = cdf2pmf_av1(default_cdf) | |
| default_rateidx = int(prob_dict[ctx_name]["init_rateidx"]) | |
| # default_rate = RATE_LIST[default_rateidx] | |
| curr_cost_list, upper_cost = self.estimate_cost_multiple( | |
| df_filtered, pmf_init=default_pmf, rateidx_init=default_rateidx | |
| ) | |
| init_cost = curr_cost_list[default_rateidx] | |
| curr_cost_list = curr_cost_list - init_cost | |
| result_dict[ctx_name] = { | |
| "current_cost": curr_cost_list, | |
| "adapt_rate": rate_parameters, | |
| "initial_cost": init_cost, | |
| "codec_cost": codec_cost, | |
| "upper_cost": upper_cost, | |
| "num_samples": len(df_filtered), | |
| "initializer": prob_dict[ctx_name]["initializer"], | |
| "init_rateidx": prob_dict[ctx_name]["init_rateidx"], | |
| "value_count": value_count_list, | |
| "value_cost": value_cost_list, | |
| } | |
| print("Training...") | |
| self.resulting_model = model.EntropyContext( | |
| ctx_group_name, num_symb, num_dims, size_list, result_dict | |
| ) | |
| self.write_results() | |
| return self.resulting_model | |
| def _get_output_filename(self): | |
| # get path and filename | |
| fullpath_data = self._filename_data | |
| filename = os.path.basename(fullpath_data) | |
| dirname = os.path.dirname(fullpath_data) | |
| if len(dirname) == 0: | |
| dirname = "." | |
| base_filename = filename.split(".")[0] | |
| fullpath_result_json = dirname + "/" + "Result_" + base_filename + ".json" | |
| return fullpath_result_json | |
| def get_searched_rate_parameter_list(self): | |
| return self._rate_search_list | |
| def write_results(self): | |
| # information | |
| ctx_group_name = self.resulting_model.ctx_group_name | |
| num_symb = self.resulting_model.num_symb | |
| num_dims = self.resulting_model.num_dims | |
| size_list = self.resulting_model.size_list | |
| result_dict = self.resulting_model.model_dict | |
| # collect keys | |
| key_list = ["information"] | |
| key_list.extend(list(result_dict.keys())) | |
| json_dict = dict.fromkeys(key_list) | |
| for key in key_list: | |
| if key == "information": | |
| json_dict[key] = { | |
| "header": { | |
| "ctx_group_name": ctx_group_name, | |
| "num_symb": num_symb, | |
| "num_dims": num_dims, | |
| "size_list": size_list, | |
| }, | |
| "training": { | |
| "rate_search_list": self._rate_search_list, | |
| "cost_regularization": self._cost_regularization, | |
| "min_sample": self._min_sample, | |
| }, | |
| } | |
| else: | |
| cost_list = result_dict[key]["current_cost"].copy() | |
| min_idx = np.argmin(cost_list) | |
| json_dict[key] = { | |
| "best_rate": self._rate_search_list[min_idx], | |
| "best_rate_idx": int(min_idx), | |
| "current_cost": cost_list.astype(int).tolist(), | |
| "initial_cost": int(result_dict[key]["initial_cost"]), | |
| "codec_cost": int(result_dict[key]["codec_cost"]), | |
| "upper_cost": int(result_dict[key]["upper_cost"]), | |
| "num_samples": int(result_dict[key]["num_samples"]), | |
| "initializer": result_dict[key]["initializer"], | |
| "init_rateidx": result_dict[key]["init_rateidx"], | |
| "value_count": result_dict[key]["value_count"], | |
| "value_cost": result_dict[key]["value_cost"], | |
| } | |
| # write to json | |
| with open(self._output_filename, "w") as jsonfile: | |
| json.dump(json_dict, jsonfile, indent=4) |