diff --git a/scripts/scaling/step1.py b/scripts/scaling/step1.py index 302d05d7d..f60ee3c78 100644 --- a/scripts/scaling/step1.py +++ b/scripts/scaling/step1.py @@ -87,6 +87,8 @@ def predict_step1(configs, data_by_name, coefficients, y_metric): predicted_data_by_name = {} plotted_predicted_data_by_name = {} + unsigned_rel_errors = [] + dmin = 0.8 * min([min(data["ds"]) for data in data_by_name.values()]) dmax = 1.2 * max([max(data["ds"]) for data in data_by_name.values()]) @@ -113,8 +115,13 @@ def predict_step1(configs, data_by_name, coefficients, y_metric): predicted_data = predicted_data_by_name[name] for d, y, y_pred in zip(data["ds"], data["ys"], predicted_data["ys"]): rel_error = (y_pred - y) / y + else: + predicted_data = predicted_data_by_name[name] + for f, y, y_pred in zip(data["fs"], data["ys"], predicted_data["ys"]): + rel_error_t = (y_pred - y) / y + unsigned_rel_errors.append(np.abs(rel_error_t)) - return predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error) + return predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error), unsigned_rel_errors def str_chinchilla_n_d_fit(coefficients): @@ -201,6 +208,8 @@ def main(): num_cols = min(4, num_tasks) num_rows = (num_tasks + num_cols - 1) // num_cols + fitting_error = 0 + if args.output_path: fig, axes = plt.subplots(num_rows, num_cols, figsize=(3.75 * num_cols, 3.25 * num_rows), squeeze=False) @@ -215,10 +224,17 @@ def main(): coefficients, cov = fit_step1(data_by_name, args.y_metric) # make predictions - predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error) = predict_step1( - configs, data_by_name, coefficients, y_metric=args.y_metric - ) - results += f"\n{task_name} | {prettify(y, False)} | {prettify(y_pred, False)} | {prettify(rel_error)}" + ( + predicted_data_by_name, + plotted_predicted_data_by_name, + (y, y_pred, rel_error), + unsigned_rel_errors, + ) = predict_step1(configs, data_by_name, coefficients, y_metric=args.y_metric) + + avg_unsigned_rel_error = np.mean(unsigned_rel_errors) + fitting_error += avg_unsigned_rel_error + + results += f"\n{task_name} | {prettify(y, False)} | {prettify(y_pred, False)} | {prettify(rel_error)} | {prettify(avg_unsigned_rel_error)}" if args.output_path: plot_step1( @@ -237,6 +253,7 @@ def main(): fig.savefig(args.output_path, dpi=300) print(results) + print("Total fitting error: ", prettify(fitting_error / num_tasks)) if __name__ == "__main__": diff --git a/scripts/scaling/step1_flops.py b/scripts/scaling/step1_flops.py index 92e8f4c75..801f0e73a 100644 --- a/scripts/scaling/step1_flops.py +++ b/scripts/scaling/step1_flops.py @@ -49,8 +49,8 @@ def fit_step1(data_by_name, y_metric): train_ys += data["ys"] if y_metric == "rc_bpb": - p0 = [-0.1, 0.3, 0.1] - bounds = ([-np.inf, -np.inf, -np.inf], [0.0, np.inf, np.inf]) + p0 = [2.0, -0.3, 0.1] + bounds = ([0, -np.inf, -np.inf], [np.inf, 0, np.inf]) coefficients, cov = get_coefficients( train_fs, train_ys, @@ -70,6 +70,8 @@ def predict_step1(configs, data_by_name, coefficients, y_metric): predicted_data_by_name = {} plotted_predicted_data_by_name = {} + unsigned_rel_errors = [] + fmin = 0.8 * min([min(data["fs"]) for data in data_by_name.values()]) fmax = 1.2 * max([max(data["fs"]) for data in data_by_name.values()]) @@ -95,8 +97,13 @@ def predict_step1(configs, data_by_name, coefficients, y_metric): predicted_data = predicted_data_by_name[name] for f, y, y_pred in zip(data["fs"], data["ys"], predicted_data["ys"]): rel_error = (y_pred - y) / y + else: + predicted_data = predicted_data_by_name[name] + for f, y, y_pred in zip(data["fs"], data["ys"], predicted_data["ys"]): + rel_error_t = (y_pred - y) / y + unsigned_rel_errors.append(np.abs(rel_error_t)) - return predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error) + return predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error), unsigned_rel_errors def str_chinchilla_flops_fit(coefficients): @@ -182,6 +189,8 @@ def main(): num_cols = min(4, num_tasks) num_rows = (num_tasks + num_cols - 1) // num_cols + fitting_error = 0 + if args.output_path: fig, axes = plt.subplots(num_rows, num_cols, figsize=(3.75 * num_cols, 3.25 * num_rows), squeeze=False) @@ -196,10 +205,17 @@ def main(): coefficients, cov = fit_step1(data_by_name, args.y_metric) # make predictions - predicted_data_by_name, plotted_predicted_data_by_name, (y, y_pred, rel_error) = predict_step1( - configs, data_by_name, coefficients, y_metric=args.y_metric - ) - results += f"\n{task_name} | {prettify(y, False)} | {prettify(y_pred, False)} | {prettify(rel_error)}" + ( + predicted_data_by_name, + plotted_predicted_data_by_name, + (y, y_pred, rel_error), + unsigned_rel_errors, + ) = predict_step1(configs, data_by_name, coefficients, y_metric=args.y_metric) + + avg_unsigned_rel_error = np.mean(unsigned_rel_errors) + fitting_error += avg_unsigned_rel_error + + results += f"\n{task_name} | {prettify(y, False)} | {prettify(y_pred, False)} | {prettify(rel_error)} | {prettify(avg_unsigned_rel_error)}" if args.output_path: plot_step1( @@ -218,6 +234,7 @@ def main(): fig.savefig(args.output_path, dpi=300) print(results) + print("Total fitting error: ", prettify(fitting_error / num_tasks)) if __name__ == "__main__":