| """ |
| 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) |
