SepMeta/run_all_models.py at master · 2003Tim/SepMeta · 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
import subprocess
import sys
import os
from datetime import datetime

PYTHON = sys.executable
THIS_DIR = os.path.dirname(__file__)
TRAIN_SCRIPT = os.path.join(THIS_DIR, "train_fsl.py")
LOG_DIR = os.path.join(THIS_DIR, "multi_run_logs")
os.makedirs(LOG_DIR, exist_ok=True)

DATASET_LIST = ["MiniImageNet", "TieredImageNet"]
INIT_WEIGHTS = {
    "MiniImageNet": "/srv/storage/hdd/zqc/wlb/code/SepMetaSyncfolders/MiniImageNet-Res12-Pre/0.1_0.1_[350, 400, 440, 460, 480]/model_best.pth.tar",
    "TieredImageNet": "/srv/storage/hdd/zqc/wlb/code/SepMetaSyncfolders/TieredImagenet-Res12-Pre/0.1_0.1_[350, 400, 440, 460, 480]/model_best.pth.tar"
}
MODEL_LIST = [
    "MatchNet", "BILSTM", "MAML",
    "DeepSet", "GCN", "FEAT", "FEATSTAR",
    "SemiFEAT", "SemiProtoFEAT"
]

CONFIG_LIST = [
    {"name": "baseline",   "extra_args": []},
    {"name": "spl",        "extra_args": ["--use_spl"]},
    {"name": "spl_cqcr",   "extra_args": ["--use_spl", "--use_cqcr"]},
]

def build_tasks():
    tasks = []
    for dataset in DATASET_LIST:
        for model in MODEL_LIST:
            for cfg in CONFIG_LIST:
                tasks.append({
                    "dataset": dataset,
                    "model": model,
                    "config": cfg["name"],
                    "extra_args": cfg["extra_args"],
                })
    return tasks

def run_task(task):
    dataset   = task["dataset"]
    model     = task["model"]
    cfg_name  = task["config"]
    extra_args = task["extra_args"]

    time_tag = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = os.path.join(LOG_DIR, f"{dataset}_{model}_{cfg_name}_{time_tag}.log")

    # model_name 用来区分保存目录，可以带上配置名
    model_name_for_save = f"{model}-{cfg_name}"
    init_path = INIT_WEIGHTS.get(dataset)

    if init_path is None:
        raise ValueError(f"No init_weights specified for dataset {dataset}, please fill INIT_WEIGHTS mapping.")
    cmd = [
        PYTHON, TRAIN_SCRIPT,
        "--dataset", dataset,
        "--model_class", model,
        "--model_name", model_name_for_save,
        "--gpu", "0",
        "--init_weights", init_path,
    ] + extra_args

    print(f"\n[RUN] ds={dataset}, model={model}, cfg={cfg_name}")
    print(f"[CMD] {' '.join(cmd)}")
    print(f"[LOG] {log_file}\n")

    with open(log_file, "w") as f:
        proc = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT)

    return proc.returncode

def main():
    tasks = build_tasks()
    failed = []
    for t in tasks:
        ret = run_task(t)
        if ret != 0:
            failed.append((t, ret))
            print(f"[FAIL] {t['dataset']} | {t['model']} | {t['config']} | code={ret}")
        else:
            print(f"[OK]   {t['dataset']} | {t['model']} | {t['config']}")

    print("\n===== SUMMARY =====")
    if failed:
        print("Failed tasks:")
        for t, code in failed:
            print(f"  {t['dataset']} | {t['model']} | {t['config']} | code={code}")
    else:
        print("All tasks finished successfully.")
    print("===================\n")

if __name__ == "__main__":
    main()