Example #1
0
def hyper_train(words_type, force_retrain, use_validation, dry_run):
    """MAKEDOC: what is hyper_train doing?"""
    logg = logging.getLogger(f"c.{__name__}.hyper_train")
    logg.debug("Start hyper_train")

    # TODO add dense_width 64, dropout 02

    ###########################################################################
    #   big grid
    ###########################################################################
    hypa_grid_big = {}

    # hypa_grid_big["base_dense_width"] = [16, 32, 64]
    # hypa_grid_big["base_dense_width"] = [16, 32, 64, 128]
    # hypa_grid_big["base_dense_width"] = [128]
    hypa_grid_big["base_dense_width"] = [32]

    # hypa_grid_big["base_filters"] = [10, 20, 30, 32, 64, 128]
    # hypa_grid_big["base_filters"] = [20, 32]
    hypa_grid_big["base_filters"] = [32]

    # hypa_grid_big["batch_size"] = [16, 32, 64]
    hypa_grid_big["batch_size"] = [32]

    ds = []
    # ds.extend(["mel01", "mel04", "mela1"])
    # ds.extend(["mel04", "mela1"])
    ds.extend(["mel04"])
    # ds.extend(["mela1"])
    # ds.extend(["mel01", "mel02", "mel03", "mel04"])
    # ds.extend(["mfcc01", "mfcc02", "mfcc03", "mfcc04"])
    # ds.extend(["aug02", "aug03", "aug04", "aug05"])
    # ds.extend(["aug06", "aug07", "aug08", "aug09"])
    # ds.extend(["aug10", "aug11", "aug12", "aug13"])
    # ds.extend(["aug14", "aug15", "aug16", "aug17"])
    # ds.extend(["aug14", "aug15", "aug16"])
    # ds.extend(["aug07"])
    # ds.extend(["aug14"])
    # ds.extend(["meL04"])
    # ds.extend(["mel05"])
    # ds.extend(["auL06", "auL07", "auL08", "auL09"])
    # ds.extend(["auL18", "auL19", "auL20", "auL21"])
    # ds.extend(["aug18", "aug19", "aug20", "aug21"])
    hypa_grid_big["dataset"] = ds

    # hypa_grid_big["dropout_type"] = ["01", "02"]
    hypa_grid_big["dropout_type"] = ["01"]

    # hypa_grid_big["epoch_num"] = [15, 30, 60]
    hypa_grid_big["epoch_num"] = [18]
    # hypa_grid_big["epoch_num"] = [15, 30]

    # hypa_grid_big["kernel_size_type"] = ["01", "02", "03"]
    hypa_grid_big["kernel_size_type"] = ["02"]

    # hypa_grid_big["pool_size_type"] = ["01", "02", "03"]
    hypa_grid_big["pool_size_type"] = ["03"]

    lr = []
    # lr.extend(["01", "02", "03"])  # fixed
    # lr.extend(["e1"])  # exp_decay_keras_01
    lr.extend(["04"])  # exp_decay_step_01
    # lr.extend(["05"])  # exp_decay_smooth_01
    # lr.extend(["06"])  # exp_decay_smooth_02
    hypa_grid_big["learning_rate_type"] = lr

    hypa_grid_big["optimizer_type"] = ["a1"]

    hypa_grid_big["words"] = [words_type]
    # hypa_grid_big["words"] = ["LTnum", "LTnumLS"]
    # hypa_grid_big["words"] = ["f2", "f1", "num", "dir", "k1", "w2", "all"]

    ###########################################################################
    #   tiny grid
    ###########################################################################
    hypa_grid_tiny = {}
    hypa_grid_tiny["base_filters"] = [20]
    hypa_grid_tiny["kernel_size_type"] = ["01"]
    hypa_grid_tiny["pool_size_type"] = ["01"]
    hypa_grid_tiny["base_dense_width"] = [32]
    hypa_grid_tiny["dropout_type"] = ["01"]
    hypa_grid_tiny["batch_size"] = [32]
    hypa_grid_tiny["epoch_num"] = [15]
    hypa_grid_tiny["learning_rate_type"] = ["01"]
    hypa_grid_tiny["optimizer_type"] = ["a1"]
    hypa_grid_tiny["dataset"] = ["mel04"]
    hypa_grid_tiny["words"] = [words_type]

    ###########################################################################
    #   best params generally
    ###########################################################################
    hypa_grid_best = {}
    hypa_grid_best["base_dense_width"] = [32]
    hypa_grid_best["base_filters"] = [20]
    hypa_grid_best["batch_size"] = [32]
    hypa_grid_best["dataset"] = ["mel01"]
    hypa_grid_best["dropout_type"] = ["01"]
    hypa_grid_best["epoch_num"] = [16]
    hypa_grid_best["kernel_size_type"] = ["02"]
    hypa_grid_best["pool_size_type"] = ["02"]
    hypa_grid_best["learning_rate_type"] = ["02"]
    hypa_grid_best["optimizer_type"] = ["a1"]
    # hypa_grid_best["words"] = ["f2", "f1", "num", "dir", "k1", "w2", "all"]
    hypa_grid_best["words"] = [words_type]

    # TODO finish kernel pool filters dense for mela1
    # hypa_grid_todo = {
    #     "base_dense_width": [16, 32, 64, 128],
    #     "base_filters": [10, 20, 32, 64, 128],
    #     "batch_size": [32],
    #     "dataset": ["mela1"],
    #     "dropout_type": ["01"],
    #     "epoch_num": [15],
    #     "kernel_size_type": ["01", "02"],
    #     "pool_size_type": ["01", "02"],
    #     "learning_rate_type": ["04", "05"],
    #     "optimizer_type": ["a1"],
    #     "words": ["f1"],
    # }

    # hypa_grid = hypa_grid_tiny
    # hypa_grid = hypa_grid_best
    hypa_grid = hypa_grid_big
    logg.debug(f"hypa_grid = {hypa_grid}")

    the_grid = list(ParameterGrid(hypa_grid))

    num_hypa = len(the_grid)
    logg.debug(f"num_hypa: {num_hypa}")

    if dry_run:
        tra_info = {"already_trained": 0, "to_train": 0}
        for hypa in the_grid:
            train_status = train_model_cnn_dry(hypa)
            tra_info[train_status] += 1
        logg.debug(f"tra_info: {tra_info}")
        return

    # check that the data is available
    for dn in hypa_grid["dataset"]:
        for wt in hypa_grid["words"]:
            logg.debug(f"\nwt: {wt} dn: {dn}\n")
            if dn.startswith("me") or dn.startswith("mfcc"):
                logg.debug(f"Preprocess! {wt} {dn}")
                preprocess_spec(dn, wt)
            elif dn.startswith("au"):
                do_augmentation(dn, wt)

    for i, hypa in enumerate(the_grid):
        logg.debug(f"\nSTARTING {i+1}/{num_hypa} with hypa: {hypa}")
        with Pool(1) as p:
            p.apply(train_model, (hypa, force_retrain))
Example #2
0
def recompute_fscore_cnn() -> None:
    """MAKEDOC: what is recompute_fscore_cnn doing?"""
    logg = logging.getLogger(f"c.{__name__}.recompute_fscore_cnn")
    # logg.setLevel("INFO")
    logg.debug("Start recompute_fscore_cnn")

    info_folder = Path("info")
    trained_folder = Path("trained_models")

    for model_folder in info_folder.iterdir():
        # logg.debug(f"model_folder: {model_folder}")

        # check that it is a CNN
        model_name = model_folder.name
        if not model_name.startswith("CNN"):
            continue

        # check that the model is trained and not a placeholder
        model_path = trained_folder / f"{model_name}.h5"
        found_model = False
        if model_path.exists():
            if model_path.stat().st_size > 100:
                found_model = True
        if not found_model:
            continue

        # load it
        model = models.load_model(model_path)

        res_recap_path = model_folder / "results_recap.json"
        if not res_recap_path.exists():
            continue
        results_recap = json.loads(res_recap_path.read_text())
        # logg.debug(f"results_recap['cm']: {results_recap['cm']}")

        recap_path = model_folder / "recap.json"
        recap = json.loads(recap_path.read_text())
        # logg.debug(f"recap['words']: {recap['words']}")

        words = recap["words"]
        hypa = recap["hypa"]

        # check that the data is available
        dn = hypa["dataset"]
        wt = hypa["words"]
        if dn.startswith("mel") or dn.startswith("mfcc"):
            preprocess_spec(dn, wt)
        elif dn.startswith("aug"):
            do_augmentation(dn, wt)

        processed_path = Path("data_proc") / f"{hypa['dataset']}"
        data, labels = load_processed(processed_path, words)

        y_pred = model.predict(data["testing"])
        cm = pred_hot_2_cm(labels["testing"], y_pred, words)
        fscore = analyze_confusion(cm, words)
        # logg.debug(f"fscore: {fscore}")

        # overwrite the cm
        results_recap["cm"] = cm.tolist()
        # add the fscore
        results_recap["fscore"] = fscore
        # increase the version
        results_recap["results_recap_version"] = "002"
        # write the new results
        res_recap_path.write_text(json.dumps(results_recap, indent=4))

        # increase the recap version (shows that it is after this debacle)
        recap["version"] = "002"
        recap_path.write_text(json.dumps(recap, indent=4))

        # save the new plots
        fig, ax = plt.subplots(figsize=(12, 12))
        plot_confusion_matrix(cm, ax, model_name, words, fscore)
        plot_cm_path = info_folder / "test_confusion_matrix.png"
        fig.savefig(plot_cm_path)
        plt.close(fig)
def hyper_train_attention(
    words_type: str,
    force_retrain: bool,
    use_validation: bool,
    dry_run: bool,
    do_find_best_lr: bool,
) -> None:
    """MAKEDOC: what is hyper_train_attention doing?"""
    logg = logging.getLogger(f"c.{__name__}.hyper_train_attention")
    # logg.setLevel("INFO")
    logg.debug("Start hyper_train_attention")

    # TODO: dropout 12 kernel 2

    # TODO fix best params then change all datasets
    # TODO auA5678 auA1234 on a bit of everything

    hypa_grid: ty.Dict[str, ty.List[str]] = {}

    ###### the words to train on
    hypa_grid["words_type"] = [words_type]

    ###### the dataset to train on
    ds = []
    # ds.extend(["mfcc01", "mfcc03", "mfcc04", "mfcc08"])
    # ds.extend(["mfcc04"])
    # ds.extend(["mel01", "mel04", "mel05", "mela1"])
    # ds.extend(["aug02", "aug03", "aug04", "aug05"])
    # ds.extend(["aug06", "aug07", "aug08", "aug09"])
    # ds.extend(["aug10", "aug11", "aug12", "aug13"])
    # ds.extend(["aug14", "aug15", "aug16", "aug17"])
    # ds.extend(["mela1"])
    ds.extend(["mel04"])
    # ds.extend(["aug15"])
    # ds.extend(["aug14"])
    # ds.extend(["aug07"])
    # hypa_grid["dataset_name"] = ["mela1"]
    # hypa_grid["dataset_name"] = ["mel04"]
    # hypa_grid["dataset_name"] = ["aug01"]
    # hypa_grid["dataset_name"] = ["aug01", "mela1", "mel04"]
    # hypa_grid["dataset_name"] = ["mela1", "mel04"]

    # ds.extend(["meL04"])
    # ds.extend(["meLa1"])
    # ds.extend(["meLa2"])
    # ds.extend(["meLa3"])
    # ds.extend(["meLa4"])
    # ds.extend(["auL18", "auL19", "auL20", "auL21"])
    # ds.extend(["auL18"])

    hypa_grid["dataset_name"] = ds

    ###### how big are the first conv layers
    # hypa_grid["conv_size_type"] = ["01", "02"]
    hypa_grid["conv_size_type"] = ["02"]

    ###### dropout after conv, 0 to skip it
    # hypa_grid["dropout_type"] = ["01", "02"]
    hypa_grid["dropout_type"] = ["01"]

    ###### the shape of the kernels in the conv layers
    # hypa_grid["kernel_size_type"] = ["01", "02"]
    hypa_grid["kernel_size_type"] = ["01"]

    ###### the dimension of the LSTM
    # hypa_grid["lstm_units_type"] = ["01", "02"]
    hypa_grid["lstm_units_type"] = ["01"]

    ###### the query style type
    qs = []
    # hypa_grid["query_style_type"] = ["01", "02", "03", "04", "05"]
    qs.extend(["01"])  # dense01
    # qs.extend(["03"])  # conv02 (LSTM)
    # qs.extend(["04"])  # conv03 (inputs)
    # qs.extend(["05"])  # dense02
    hypa_grid["query_style_type"] = qs

    ###### the width of the dense layers
    # hypa_grid["dense_width_type"] = ["01", "02"]
    hypa_grid["dense_width_type"] = ["01"]

    ###### the learning rates for the optimizer
    lr = []
    # lr.extend(["01", "02"])  # fixed
    # lr.extend(["01"])  # fixed
    # lr.extend(["02"])  # fixed
    lr.extend(["03"])  # exp_decay_step_01
    # lr.extend(["04"])  # exp_decay_smooth_01
    # lr.extend(["05"])  # clr_triangular2_01
    # lr.extend(["06"])  # clr_triangular2_02
    # lr.extend(["07"])  # clr_triangular2_03
    # lr.extend(["08"])  # clr_triangular2_04
    # lr.extend(["09"])  # clr_triangular2_05
    # lr.extend(["10"])  # exp_decay_smooth_02

    hypa_grid["learning_rate_type"] = lr

    ###### which optimizer to use
    # hypa_grid["optimizer_type"] = ["a1", "r1"]
    hypa_grid["optimizer_type"] = ["a1"]

    ###### the batch size to use
    # hypa_grid["batch_size_type"] = ["01", "02"]
    hypa_grid["batch_size_type"] = ["02"]

    ###### the number of epochs
    en = []
    # en.extend(["01"])  # 15
    en.extend(["02"])  # 30
    # en.extend(["03", "04"])
    hypa_grid["epoch_num_type"] = en

    # the grid you are generating from (useful to recreate the training)
    logg.debug(f"hypa_grid: {hypa_grid}")

    # create the list of combinations
    the_grid = list(ParameterGrid(hypa_grid))

    # from retrain_list import retrain_grid
    # the_grid = retrain_grid
    # for hypa in the_grid:
    #     dn = hypa["dataset_name"]
    #     wt = hypa["words_type"]
    #     logg.debug(f"\nwt: {wt} dn: {dn}\n")
    #     if dn.startswith("me"):
    #         preprocess_spec(dn, wt)
    #     elif dn.startswith("au"):
    #         do_augmentation(dn, wt)

    # MAYBE finish this on yn
    # hypa_grid: {'words_type': ['yn'], 'dataset_name': ['mela1', 'mel04', 'aug15',
    # 'aug14', 'aug07'], 'conv_size_type': ['02'], 'dropout_type': ['01'],
    # 'kernel_size_type': ['01', '02'], 'lstm_units_type': ['01'], 'query_style_type':
    # ['03', '04', '05'], 'dense_width_type': ['01'], 'learning_rate_type': ['03',
    # '04'], 'optimizer_type': ['a1'], 'batch_size_type': ['02'], 'epoch_num_type':
    # ['01']}

    # hypa_base = {
    #     "batch_size_type": "02",
    #     "conv_size_type": "02",
    #     # "dataset_name": "mel04L",
    #     # "dataset_name": "mel04",
    #     "dense_width_type": "01",
    #     "dropout_type": "01",
    #     "epoch_num_type": "02",
    #     "kernel_size_type": "01",
    #     "learning_rate_type": "03",
    #     "lstm_units_type": "01",
    #     "optimizer_type": "a1",
    #     "query_style_type": "01",
    #     # "words_type": "LTnumLS",
    #     # "words_type": "LTBnumLS",
    #     # "words_type": "LTnum",
    # }
    # the_grid = []

    # from copy import deepcopy
    # hypa = deepcopy(hypa_base)
    # hypa["dataset_name"] = "mel04"
    # hypa["words_type"] = "LTBnum"
    # the_grid.append(hypa)
    # hypa = deepcopy(hypa_base)
    # hypa["dataset_name"] = "mel04"
    # hypa["words_type"] = "LTBall"
    # the_grid.append(hypa)

    # hypa = deepcopy(hypa_base)
    # hypa["dataset_name"] = "mel04L"
    # hypa["words_type"] = "LTBnumLS"
    # the_grid.append(hypa)
    # hypa = deepcopy(hypa_base)
    # hypa["dataset_name"] = "mel04L"
    # hypa["words_type"] = "LTBallLS"
    # the_grid.append(hypa)

    num_hypa = len(the_grid)
    logg.debug(f"num_hypa: {num_hypa}")

    # check which models need to be trained
    if dry_run:
        tra_info = {"already_trained": 0, "to_train": 0}
        for hypa in the_grid:
            train_status = train_model_att_dry(hypa, use_validation)
            tra_info[train_status] += 1
        logg.debug(f"tra_info: {tra_info}")
        return

    # magic to fix the GPUs
    # setup_gpus()

    # check that the data is available
    for dn in hypa_grid["dataset_name"]:
        for wt in hypa_grid["words_type"]:
            logg.debug(f"\nwt: {wt} dn: {dn}\n")
            if dn.startswith("me"):
                preprocess_spec(dn, wt)
            elif dn.startswith("au"):
                do_augmentation(dn, wt)

    # useful to pick good values of lr
    if do_find_best_lr:
        hypa = the_grid[0]
        logg.debug(f"\nSTARTING find_best_lr with hypa: {hypa}")
        find_best_lr(hypa)
        return

    # train each combination
    for i, hypa in enumerate(the_grid):
        logg.debug(f"\nSTARTING {i+1}/{num_hypa} with hypa: {hypa}")
        with Pool(1) as p:
            p.apply(train_attention, (hypa, force_retrain, use_validation))
Example #4
0
def visualize_datasets(word_index):
    """MAKEDOC: what is visualize_datasets doing?"""
    logg = logging.getLogger(f"c.{__name__}.visualize_datasets")
    logg.debug("Start visualize_datasets")

    # magic to fix the GPUs
    setup_gpus()

    # show different datasets

    # datasets = [ "mfcc01", "mfcc02", "mfcc03", "mfcc04", "mfcc05", "mfcc06", "mfcc07", "mfcc08"]
    # datasets = [ "mel01", "mel02", "mel03", "mel04", "mel05", "mel06", "mel07", "mel08",
    #     "mel09", "mel10", "mel11", "mel12", "mel13", "mel14", "mel15", "melc1", "melc2",
    #     "melc3", "melc4", "mela1", "meL04", "meLa1", "auL18", "aug18", ]
    # datasets = [ "mel01", "mel04", "mel06", "melc1" ]
    # datasets = ["mel09", "mel10", "mel11", "melc1"]
    datasets = ["mel04", "mel04a", "mel04b", "melc1"]

    # words = words_types["f1"]
    # a_word = words[0]
    # a_word = "loudest_one"
    a_word = "happy"
    # a_word = "_other_ltts_loud"

    # datasets = []
    # datasets.extend(["meL04", "meLa1", "meLa2", "meLa3", "meLa4"])
    # datasets.extend(["auL06", "auL07", "auL08", "auL09"])
    # datasets.extend(["auL18", "auL19", "auL20", "auL21"])
    # a_word = "loudest_two"

    # datasets = []
    # datasets.extend(["mel04", "mela1"])
    # datasets.extend(["aug14", "aug15"])
    # a_word = "forward"
    # datasets.extend(["aug14", "aug07"])
    # a_word = "one"
    # a_word = "_other_ltts"

    # which word in the dataset to plot
    iw = word_index

    processed_folder = Path("data_proc")

    # fig, axes = plt.subplots(4, 5, figsize=(12, 15))
    nrows, ncols = find_rowcol(len(datasets))
    base_figsize = 5
    figsize = (ncols * base_figsize * 1.5, nrows * base_figsize)
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
    if nrows * ncols > 1:
        axes_flat = axes.flat
    else:
        axes_flat = [axes]

    fig.suptitle(f"Various spectrograms for {a_word}", fontsize=20)
    for i, ax in enumerate(axes_flat[: len(datasets)]):

        # the current dataset being plotted
        dataset_name = datasets[i]
        processed_path = processed_folder / f"{dataset_name}"
        word_path = processed_path / f"{a_word}_training.npy"
        logg.debug(f"word_path: {word_path}")

        # FIXME this is shaky as hell
        if not word_path.exists():
            if dataset_name.startswith("me"):
                preprocess_spec(dataset_name, f"_{a_word}")
            elif dataset_name.startswith("au"):
                do_augmentation(dataset_name, f"_{a_word}")

        word_data = np.load(word_path, allow_pickle=True)
        logg.debug(f"{dataset_name} {a_word} shape: {word_data[iw].shape}")
        title = f"{dataset_name}: shape {word_data[iw].shape}"
        plot_spec(word_data[iw], ax, title=title)
    fig.tight_layout()

    plot_folder = Path("plot_models")
    dt_names = "_".join(datasets)
    fig.savefig(plot_folder / f"{a_word}_{dt_names}_specs.pdf")
Example #5
0
def evaluate_model_cm(model_name: str, test_words_type: str) -> float:
    r"""MAKEDOC: what is evaluate_model_cm doing?"""
    logg = logging.getLogger(f"c.{__name__}.evaluate_model_cm")
    # logg.setLevel("INFO")
    # logg.debug("\nStart evaluate_model_cm")

    # magic to fix the GPUs
    setup_gpus()

    logg.debug(f"\nmodel_name: {model_name}")

    dataset_re = re.compile("_ds(.*?)_")
    match = dataset_re.search(model_name)
    if match is not None:
        logg.debug(f"match[1]: {match[1]}")
        dataset_name = match[1]

    train_words_type_re = re.compile("_w(.*?)[_.]")
    match = train_words_type_re.search(model_name)
    if match is not None:
        logg.debug(f"match[1]: {match[1]}")
        train_words_type = match[1]

    arch_type = model_name[:3]

    if arch_type == "ATT":
        train_type_tag = "attention"
    else:
        train_type_tag = "area"

    # load the model
    model_folder = Path("trained_models") / train_type_tag
    model_path = model_folder / f"{model_name}.h5"
    model = tf_models.load_model(model_path)
    # model.summary()

    train_words = words_types[train_words_type]
    logg.debug(f"train_words: {train_words}")
    test_words = words_types[test_words_type]
    logg.debug(f"test_words: {test_words}")

    # input data must exist
    if dataset_name.startswith("mel"):
        preprocess_spec(dataset_name, test_words_type)
    elif dataset_name.startswith("aug"):
        do_augmentation(dataset_name, test_words_type)

    # input data
    processed_path = Path("data_proc") / f"{dataset_name}"
    data, labels = load_processed(processed_path, test_words)
    logg.debug(f"list(data.keys()): {list(data.keys())}")
    logg.debug(f"data['testing'].shape: {data['testing'].shape}")

    # evaluate on the words you trained on
    logg.debug("Evaluate on test data:")
    model.evaluate(data["testing"], labels["testing"])
    # model.evaluate(data["validation"], labels["validation"])

    # predict labels/cm/fscore
    y_pred = model.predict(data["testing"])
    cm = pred_hot_2_cm(labels["testing"], y_pred, test_words)
    # y_pred = model.predict(data["validation"])
    # cm = pred_hot_2_cm(labels["validation"], y_pred, test_words)
    fscore = analyze_confusion(cm, test_words)
    logg.debug(f"fscore: {fscore}")

    fig, ax = plt.subplots(figsize=(12, 12))
    plot_confusion_matrix(cm, ax, model_name, test_words, fscore, train_words)

    fig_name = f"{model_name}_test{test_words_type}_cm.{{}}"
    cm_folder = Path("plot_results") / "cm_all01"
    if not cm_folder.exists():
        cm_folder.mkdir(parents=True, exist_ok=True)

    plot_cm_path = cm_folder / fig_name.format("png")
    fig.savefig(plot_cm_path)
    plot_cm_path = cm_folder / fig_name.format("pdf")
    fig.savefig(plot_cm_path)

    # plt.show()
    return fscore
Example #6
0
def hyper_train_area(
    words_type: str,
    force_retrain: bool,
    use_validation: bool,
    dry_run: bool,
    do_find_best_lr: bool,
) -> None:
    """MAKEDOC: what is hyper_train_area doing?"""
    logg = logging.getLogger(f"c.{__name__}.hyper_train_area")
    # logg.setLevel("INFO")
    logg.debug("Start hyper_train_area")

    ##########################################################
    #   Hyper-parameters grid
    ##########################################################

    hypa_grid: ty.Dict[str, ty.List[str]] = {}

    ###### the net type
    nt = []
    nt.append("SIM")
    nt.append("SI2")
    nt.append("AAN")
    nt.append("VAN")
    hypa_grid["net_type"] = nt

    ###### the words to train on
    wtl = []
    # wtl.extend([words_type])
    # wtl.extend(["LTnum", "LTall", "yn"])
    # wtl.extend(["LTnum", "yn"])
    # wtl.extend(["LTnum", "yn", "f1", "k1"])
    # wtl.extend(["LTBnum", "LTBall"])
    # wtl.extend(["LTBnum"])
    # wtl.extend(["num"])
    # wtl.extend(["LTBnumLS", "LTBallLS"])
    wtl.extend(["LTBnumLS"])
    # wtl.extend(["lr"])
    # wtl.extend(["similar", "all"])
    # wtl.extend(["all"])

    hypa_grid["words_type"] = wtl

    ###### the dataset to train on
    ds = []

    # TODO VAN on LTall
    # TODO AAN/SIM/VAN on LTnum
    # TODO AAN on LTnum for all datasets, only one lr
    # ds.extend(["mel04"])
    # ds.extend(["mela1"])
    # ds.extend(["aug07"])
    # ds.extend(["aug14"])
    # ds.extend(["aug15"])

    # TODO auL6789 auL18901 on all net_type
    # TODO auA5678 on VAN (on LTnumLS)
    # TODO auA5678 with lr04 (on LTnumLS to complete lr03 is done)
    # MAYBE start removing ARN, too much time
    # ds.extend(["auA01", "auA02", "auA03", "auA04"])
    # ds.extend(["auA05", "auA06", "auA07", "auA08"])
    # ds.extend(["auA04"])
    ds.extend(["mel04L"])

    # TODO just the 3 best per architecture on noval

    hypa_grid["dataset_name"] = ds

    ###### the learning rates for the optimizer
    lr = []
    # lr.extend(["01", "02"])  # fixed
    lr.extend(["03"])  # exp_decay_step_01
    lr.extend(["04"])  # exp_decay_smooth_01
    # lr.extend(["05"])  # clr_triangular2_01
    # lr.extend(["06"])  # clr_triangular2_02
    hypa_grid["learning_rate_type"] = lr

    ###### which optimizer to use
    # hypa_grid["optimizer_type"] = ["a1", "r1"]
    hypa_grid["optimizer_type"] = ["a1"]

    ###### the batch size (the key is converted to int)
    bs = []
    bs.extend(["32"])
    # bs.extend(["16"])
    hypa_grid["batch_size_type"] = bs

    ###### the number of epochs (the key is converted to int)
    en = []
    en.extend(["15"])
    # en.extend(["13"])
    # en.extend(["10"])
    hypa_grid["epoch_num_type"] = en

    ###### build the combinations
    logg.debug(f"hypa_grid = {hypa_grid}")
    the_grid = list(ParameterGrid(hypa_grid))

    # hijack the grid
    # hypa = {
    #     "batch_size_type": "32",
    #     "dataset_name": "auA07",
    #     "epoch_num_type": "15",
    #     "learning_rate_type": "03",
    #     "net_type": "VAN",
    #     "optimizer_type": "a1",
    #     "words_type": "LTnumLS",
    # }
    # the_grid = [hypa]

    num_hypa = len(the_grid)
    logg.debug(f"num_hypa: {num_hypa}")

    ##########################################################
    #   Setup pre train
    ##########################################################

    train_type_tag = "area"

    # where to save the trained models
    trained_folder = Path("trained_models") / train_type_tag
    if not trained_folder.exists():
        trained_folder.mkdir(parents=True, exist_ok=True)

    # where to put the info folders
    root_info_folder = Path("info") / train_type_tag
    if not root_info_folder.exists():
        root_info_folder.mkdir(parents=True, exist_ok=True)

    # count how many models are left to train
    if dry_run:
        tra_info = {"already_trained": 0, "to_train": 0}
        for hypa in the_grid:
            train_status = train_model_area_dry(hypa, use_validation,
                                                trained_folder)
            tra_info[train_status] += 1
        logg.debug(f"tra_info: {tra_info}")
        return

    # check that the data is available
    for dn in hypa_grid["dataset_name"]:
        for wt in hypa_grid["words_type"]:
            logg.debug(f"\nwt: {wt} dn: {dn}\n")
            if dn.startswith("mel"):
                preprocess_spec(dn, wt)
            elif dn.startswith("aug"):
                do_augmentation(dn, wt)

    # useful to pick good values of lr
    if do_find_best_lr:
        hypa = the_grid[0]
        logg.debug(f"\nSTARTING find_best_lr with hypa: {hypa}")
        find_best_lr(hypa)
        return

    ##########################################################
    #   Train all hypas
    ##########################################################

    for i, hypa in enumerate(the_grid):
        logg.debug(f"\nSTARTING {i+1}/{num_hypa} with hypa: {hypa}")
        with Pool(1) as p:
            p.apply(
                train_area,
                (hypa, force_retrain, use_validation, trained_folder,
                 root_info_folder),
            )