BioLab-AAC2025-Detection/run_training_aac_grid_search.py at main · MI-BioLab/BioLab-AAC2025-Detection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import subprocess
import itertools
import os


GENERAL_PARAMETERS = [
    "--data.num_workers",
    "10",
]


def run_training(params):
    model_name_base = "dinov2_vits14"
    experiment_id = f"{model_name_base}_grid_multiclass_epoch{params['epochs']}_lr{params['lr']}_aug{params['augmentation_type']}_sched{params['scheduler_type']}"
    expected_experiment_folder = f"experiments_data/experiment_{experiment_id}/window_0"
    training_done_marker = os.path.join(expected_experiment_folder, "done")
    testing_done_marker = os.path.join(
        expected_experiment_folder, "predictions_test_all.npz"
    )

    # Build the command according to the parameters
    command = [
        "python",
        "lightning_main.py",
        "fit",
        "--config",
        "aac_training_configurations/aac_full_training.yaml",
        "--config",
        "aac_training_configurations/mono_step/monostep_multiclass.yaml",
        "--experiment_info.experiment_id",
        experiment_id,
        "--model.model_name",
        f"{model_name_base}_tune",
        "--model.model_input_size",
        "112",
        "--model.training_cropping_strategy",
        "as_is",
        "--model.evaluation_cropping_strategy",
        "as_is",
        "--data.dataset_loader.dataset_path",
        "/home/lorenzo/Desktop/challenge_data/track1_arrow",
        "--trainer.max_epochs",
        str(params["epochs"]),
        "--model.optimizer.lr",
        str(params["lr"]),
        "--model.augmentation_type",
        params["augmentation_type"],
        "--model.scheduler",
        params["scheduler_type"],
        *GENERAL_PARAMETERS,
    ]

    command_test = list(command)
    command_test[2] = "test"
    env = os.environ.copy()

    if os.path.exists(training_done_marker):
        print(f"[INFO] Experiment {experiment_id} already completed. Skipping.")
    else:
        print("Running fit command:", " ".join(command))
        subprocess.run(command, check=True, env=env)

    if os.path.exists(testing_done_marker):
        print(
            f"[INFO] Testing for experiment {experiment_id} already completed. Skipping."
        )
    else:
        print("Running test command:", " ".join(command))
        subprocess.run(command_test, check=True, env=env)


def main():
    # Grid of parameters to search over
    epochs_list = [
        32,
    ]
    lr_list = [0.01, 0.005, 0.001, 0.0005]
    augmentation_types = ["minimal", "ultra_minimal"]
    scheduler_types = ["OneCycleLR", "OneCycleLRSoft", "ReduceLROnPlateau"]

    # Iterate over all combinations using itertools.product
    for epochs, lr, aug, sched in itertools.product(
        epochs_list, lr_list, augmentation_types, scheduler_types
    ):
        params = {
            "epochs": epochs,
            "lr": lr,
            "augmentation_type": aug,
            "scheduler_type": sched,
        }
        print(f"\n[INFO] Starting grid search with params: {params}")
        try:
            run_training(params)
        except subprocess.CalledProcessError as e:
            print(
                f"[ERROR] Command failed with error: {e}. Continuing to next parameter set."
            )


if __name__ == "__main__":
    main()