diff --git a/src/scripts/AleatoricInits.py b/src/scripts/AleatoricInits.py index b81f832..4d3372d 100644 --- a/src/scripts/AleatoricInits.py +++ b/src/scripts/AleatoricInits.py @@ -291,7 +291,7 @@ def beta_type(value): ax.set_title("Deep Evidential Regression") elif model[0:2] == "DE": ax.set_title("Deep Ensemble (100 models)") - ax.set_ylim([0, 14]) + ax.set_ylim([0, 6]) plt.legend() if config.get_item("analysis", "savefig", "Analysis"): plt.savefig( diff --git a/src/scripts/LossFunctions.py b/src/scripts/LossFunctions.py index 06ab552..abb7213 100644 --- a/src/scripts/LossFunctions.py +++ b/src/scripts/LossFunctions.py @@ -3,10 +3,9 @@ import argparse import torch import matplotlib.pyplot as plt -from data import DataModules from data.data import DataPreparation from utils.config import Config -from utils.defaults import DefaultsAnalysis, DefaultsDE +from utils.defaults import DefaultsAnalysis from analyze.analyze import AggregateCheckpoints # from plots import Plots @@ -22,50 +21,50 @@ def parse_args(): # option to pass name of config parser.add_argument("--config", "-c", default=None) - # data info - parser.add_argument( - "--data_path", - "-d", - default=DefaultsAnalysis["data"]["data_path"], - ) - parser.add_argument( - "--data_engine", - "-dl", - default=DefaultsAnalysis["data"]["data_engine"], - choices=DataModules.keys(), - ) - # model + # we need some info about the model to run this analysis # path to save the model results - parser.add_argument("--out_dir", - default=DefaultsAnalysis["common"]["out_dir"]) - + parser.add_argument("--dir", default=DefaultsAnalysis["common"]["dir"]) # now args for model parser.add_argument( "--n_models", type=int, - default=DefaultsDE["model"]["n_models"], + default=DefaultsAnalysis["model"]["n_models"], help="Number of MVEs in the ensemble", ) parser.add_argument( "--BETA", type=beta_type, required=False, - default=DefaultsDE["model"]["BETA"], + default=DefaultsAnalysis["model"]["BETA"], help="If loss_type is bnn_loss, specify a beta as a float or \ there are string options: linear_decrease, \ step_decrease_to_0.5, and step_decrease_to_1.0", ) parser.add_argument( - "--noise_level_list", + "--COEFF", + type=float, + required=False, + default=DefaultsAnalysis["model"]["COEFF"], + help="COEFF for DER", + ) + parser.add_argument( + "--loss_type", type=str, required=False, + default=DefaultsAnalysis["model"]["loss_type"], + help="loss_type for DER, either SDER or DER", + ) + parser.add_argument( + "--noise_level_list", + type=list, + required=False, default=DefaultsAnalysis["analysis"]["noise_level_list"], help="Noise levels to compare", ) parser.add_argument( "--model_names_list", - type=str, + type=list, required=False, default=DefaultsAnalysis["analysis"]["model_names_list"], help="Beginning of name for saved checkpoints and figures", @@ -74,7 +73,7 @@ def parse_args(): "--n_epochs", type=int, required=False, - default=DefaultsAnalysis["analysis"]["n_epochs"], + default=DefaultsAnalysis["model"]["n_epochs"], help="number of epochs", ) parser.add_argument( @@ -117,18 +116,18 @@ def parse_args(): os.makedirs(os.path.dirname(temp_config), exist_ok=True) # check if args were specified in cli - # if not, default is from DefaultsDE dictionary input_yaml = { - "common": {"out_dir": args.out_dir}, - "data": { - "data_path": args.data_path, - "data_engine": args.data_engine, + "common": {"dir": args.dir}, + "model": { + "n_models": args.n_models, + "n_epochs": args.n_epochs, + "BETA": args.BETA, + "COEFF": args.COEFF, + "loss_type": args.loss_type, }, - "model": {"n_models": args.n_models, "BETA": args.BETA}, "analysis": { "noise_level_list": args.noise_level_list, "model_names_list": args.model_names_list, - "n_epochs": args.n_epochs, "plot": args.plot, "savefig": args.savefig, "verbose": args.verbose, @@ -165,26 +164,37 @@ def beta_type(value): DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") noise_list = config.get_item("analysis", "noise_level_list", "Analysis") color_list = config.get_item("plots", "color_list", "Analysis") + BETA = config.get_item("model", "BETA", "Analysis") + COEFF = config.get_item("model", "COEFF", "Analysis") + loss_type = config.get_item("model", "loss_type", "Analysis") sigma_list = [] for noise in noise_list: sigma_list.append(DataPreparation.get_sigma(noise)) - path_to_chk = config.get_item("common", "out_dir", "Analysis") + root_dir = config.get_item("common", "dir", "Analysis") + path_to_chk = root_dir + "checkpoints/" + path_to_out = root_dir + "analysis/" + # check that this exists and if not make it + if not os.path.isdir(path_to_out): + print("does not exist, making dir", path_to_out) + os.mkdir(path_to_out) + else: + print("already exists", path_to_out) model_name_list = config.get_item("analysis", "model_names_list", "Analysis") print("model list", model_name_list) print("noise list", noise_list) - loss = { + chk_module = AggregateCheckpoints() + mse_loss = { model_name: {noise: [] for noise in noise_list} for model_name in model_name_list } - mse_loss = { + loss = { model_name: {noise: [] for noise in noise_list} for model_name in model_name_list } - chk_module = AggregateCheckpoints() - n_epochs = config.get_item("analysis", "n_epochs", "Analysis") - n_models = config.get_item("model", "n_models", "DE") + n_epochs = config.get_item("model", "n_epochs", "Analysis") + n_models = config.get_item("model", "n_models", "Analysis") for model in model_name_list: for noise in noise_list: # now run the analysis on the resulting checkpoints @@ -194,34 +204,32 @@ def beta_type(value): model, noise, epoch, - config.get_item("model", "BETA", "DE"), DEVICE, + path=path_to_chk, + COEFF=COEFF, + loss=loss_type, ) # path=path_to_chk) # things to grab: 'valid_mse' and 'valid_bnll' mse_loss[model][noise].append(chk["valid_mse"]) loss[model][noise].append(chk["valid_loss"]) - elif model[0:2] == "DE": + elif model[0:3] == "DE_": for nmodel in range(n_models): mse_loss_one_model = [] loss_one_model = [] for epoch in range(n_epochs): - try: - chk = chk_module.load_checkpoint( - model, - noise, - epoch, - config.get_item("model", "BETA", "DE"), - DEVICE, - nmodel=nmodel, - ) - except FileNotFoundError: - continue - # path=path_to_chk) - # things to grab: 'valid_mse' and 'valid_bnll' - # print(chk) + chk = chk_module.load_checkpoint( + model, + noise, + epoch, + DEVICE, + path=path_to_chk, + BETA=BETA, + nmodel=nmodel, + ) mse_loss_one_model.append(chk["valid_mse"]) loss_one_model.append(chk["valid_loss"]) + mse_loss[model][noise].append(mse_loss_one_model) loss[model][noise].append(loss_one_model) # make a two-paneled plot for the different noise levels @@ -236,26 +244,78 @@ def beta_type(value): ax.set_title(model) # Set title for each subplot for i, noise in enumerate(noise_list): if model[0:3] == "DER": - im = ax.scatter( + ax.plot( range(n_epochs), mse_loss[model][noise], color=color_list[i], - edgecolors="black", - label=r"$\sigma = $"+str(sigma_list[i]), + label=r"$\sigma = $" + str(sigma_list[i]), ) - ax.set_title("Deep Evidential Regression") - ax.legend(im) - elif model[0:2] == "DE": + else: for n in range(n_models): - ax.scatter( + ax.plot( range(n_epochs), mse_loss[model][noise][n], - color=color_list[i], - edgecolors="black", + color=color_list[i] ) - ax.set_title("Deep Ensemble (100 models)") - ax.set_ylabel("MSE") + ax.set_ylabel("MSE Loss") ax.set_xlabel("Epoch") + if model[0:3] == "DER": + ax.set_title("Deep Evidential Regression") + plt.legend() + elif model[0:2] == "DE": + ax.set_title("Deep Ensemble (100 models)") ax.set_ylim([0, 250]) + if config.get_item("analysis", "savefig", "Analysis"): + plt.savefig( + str(path_to_out) + + "mse_loss_n_epochs_" + + str(n_epochs) + + "_n_models_DE_" + + str(n_models) + + ".png" + ) + if config.get_item("analysis", "plot", "Analysis"): + plt.show() - plt.show() + plt.clf() + fig = plt.figure(figsize=(10, 4)) + # try this instead with a fill_between method + for i, model in enumerate(model_name_list): + ax = fig.add_subplot(1, len(model_name_list), i + 1) + # Your plotting code for each model here + ax.set_title(model) # Set title for each subplot + for i, noise in enumerate(noise_list): + if model[0:3] == "DER": + ax.plot( + range(n_epochs), + loss[model][noise], + color=color_list[i], + label=r"$\sigma = $" + str(sigma_list[i]), + ) + else: + for n in range(n_models): + ax.plot( + range(n_epochs), + loss[model][noise][n], + color=color_list[i], + ) + ax.set_xlabel("Epoch") + if model[0:3] == "DER": + ax.set_title("Deep Evidential Regression") + ax.set_ylabel("NIG Loss") + plt.legend() + elif model[0:3] == "DE_": + ax.set_title("Deep Ensemble (100 models)") + ax.set_ylabel(r"$\beta-$NLL Loss") + # ax.set_ylim([0, 11]) + if config.get_item("analysis", "savefig", "Analysis"): + plt.savefig( + str(path_to_out) + + "loss_n_epochs_" + + str(n_epochs) + + "_n_models_DE_" + + str(n_models) + + ".png" + ) + if config.get_item("analysis", "plot", "Analysis"): + plt.show() diff --git a/src/train/train.py b/src/train/train.py index 11adf30..e579267 100644 --- a/src/train/train.py +++ b/src/train/train.py @@ -132,12 +132,6 @@ def train_DER( zorder=100, ls="None", ) - ax1.annotate( - r"med $u_{ep} = $" + str(np.median(loss[2])), - xy=(0.03, 0.93), - xycoords="axes fraction", - color="black", - ) """ else: ax1.scatter(y, diff --git a/src/utils/defaults.py b/src/utils/defaults.py index 21d6de4..fe52a4e 100644 --- a/src/utils/defaults.py +++ b/src/utils/defaults.py @@ -60,9 +60,9 @@ # the engines are the classes, defined "model_engine": "DER", "model_type": "DER", - "loss_type": "SDER", + "loss_type": "DER", "init_lr": 0.001, - "COEFF": 0.5, + "COEFF": 0.01, "n_epochs": 100, "save_all_checkpoints": False, "save_final_checkpoint": False, @@ -95,8 +95,8 @@ "n_models": 100, "n_epochs": 100, "BETA": 0.5, - "COEFF": 0.5, - "loss_type": "SDER" + "COEFF": 0.01, + "loss_type": "DER" }, "analysis": { "noise_level_list": ["low", "medium", "high"],