Upload scripts/run_multiseed_baseline.py with huggingface_hub

scripts/run_multiseed_baseline.py

@@ -0,0 +1,203 @@
+"""
+run_multiseed_baseline.py
+Multi-seed training for DLoc baseline (Enc_2Dec_Network).
+"""
+import os, sys, time, random
+import torch
+import numpy as np
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=FutureWarning)
+
+from utils import localization_error, write_log
+from data_loader import load_data
+from modelADT import ModelADT
+from joint_model import Enc_2Dec_Network
+
+SEEDS = [42, 123, 777]
+EPOCHS = 120
+BATCH_SIZE = 64
+NUM_WORKERS = 4
+
+trainpath = [
+    "./data/dataset_jacobs_July28.mat",
+    "./data/dataset_non_fov_train_jacobs_July28_2.mat",
+    "./data/dataset_fov_train_jacobs_July28_2.mat",
+]
+testpath = [
+    "./data/dataset_fov_test_jacobs_July28_2.mat",
+    "./data/dataset_non_fov_test_jacobs_July28_2.mat",
+]
+
+SEP1 = "=" * 60
+SEP2 = "=" * 65
+SEP3 = "-" * 65
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def build_model():
+    from params_storage.params_fig10b_single_gpu import opt_exp, opt_encoder, opt_decoder, opt_offset_decoder
+    enc = ModelADT(); enc.initialize(opt_encoder); enc.setup(opt_encoder)
+    dec = ModelADT(); dec.initialize(opt_decoder); dec.setup(opt_decoder)
+    off = ModelADT(); off.initialize(opt_offset_decoder); off.setup(opt_offset_decoder)
+    model = Enc_2Dec_Network()
+    model.initialize(opt_exp, enc, dec, off, gpu_ids=opt_exp.gpu_ids)
+    return model
+
+
+def train_one_seed(seed, train_loader, test_loader, save_dir):
+    set_seed(seed)
+    print(f"\n{SEP1}")
+    print(f"  SEED {seed} -- BASELINE DLoc ({EPOCHS} epochs)")
+    print(SEP1)
+
+    model = build_model()
+    seed_dir = os.path.join(save_dir, f"seed_{seed}")
+    os.makedirs(seed_dir, exist_ok=True)
+
+    best_median = float("inf")
+    best_p90 = float("inf")
+    best_p99 = float("inf")
+
+    for epoch in range(1, EPOCHS + 1):
+        t0 = time.time()
+
+        # TRAIN
+        model.encoder.net.train()
+        model.decoder.net.train()
+        model.offset_decoder.net.train()
+        for data in train_loader:
+            model.set_input(data[1], data[2], data[0], shuffle_channel=True)
+            model.optimize_parameters()
+
+        # LR scheduling (critical — linear decay from niter to niter+niter_decay)
+        model.encoder.update_learning_rate()
+        model.decoder.update_learning_rate()
+        model.offset_decoder.update_learning_rate()
+
+        # TEST
+        model.eval()
+        test_errors = []
+        with torch.no_grad():
+            for data in test_loader:
+                model.set_input(data[1], data[2], data[0], shuffle_channel=False)
+                model.test()
+                pred = model.decoder.output.cpu().numpy()
+                lbl = data[2].cpu().numpy()
+                if lbl.ndim == 3:
+                    lbl = lbl[:, np.newaxis, :, :]
+                test_errors.extend(localization_error(pred, lbl, scale=0.1))
+
+        test_median = np.median(test_errors)
+        test_p90 = np.percentile(test_errors, 90)
+        test_p99 = np.percentile(test_errors, 99)
+        elapsed = time.time() - t0
+
+        rl = f"seed{seed}_baseline"
+        write_log([str(test_median)], rl, log_dir=seed_dir, log_type="test_median_error")
+        write_log([str(test_p90)], rl, log_dir=seed_dir, log_type="test_90_error")
+        write_log([str(test_p99)], rl, log_dir=seed_dir, log_type="test_99_error")
+
+        marker = ""
+        if test_median < best_median:
+            best_median = test_median
+            best_p90 = test_p90
+            best_p99 = test_p99
+            marker = " *"
+
+        if epoch % 10 == 0 or epoch == 1 or marker:
+            print(f"  Epoch {epoch}/{EPOCHS} | {elapsed:.0f}s | "
+                  f"med={test_median:.3f}m P90={test_p90:.3f}m P99={test_p99:.3f}m{marker}")
+
+    print(f"  Seed {seed} done. Best: med={best_median:.4f}m P90={best_p90:.3f}m P99={best_p99:.3f}m")
+    return {"median": best_median, "p90": best_p90, "p99": best_p99}
+
+
+def main():
+    device = torch.device("cuda:0")
+    print(f"Device: {device}")
+    print(f"Seeds: {SEEDS} | Epochs: {EPOCHS}")
+
+    print("\nLoading training data...")
+    train_wo, train_w, train_lbl = [], [], []
+    for fp in trainpath:
+        wo, w, lbl = load_data(fp)
+        if lbl.dim() == 3:
+            lbl = lbl.unsqueeze(1)
+        train_wo.append(wo)
+        train_w.append(w)
+        train_lbl.append(lbl)
+    train_wo = torch.cat(train_wo, 0)
+    train_w = torch.cat(train_w, 0)
+    train_lbl = torch.cat(train_lbl, 0)
+    print(f"Training samples: {train_w.shape[0]}")
+
+    print("Loading testing data...")
+    test_wo, test_w, test_lbl = [], [], []
+    for fp in testpath:
+        wo, w, lbl = load_data(fp)
+        if lbl.dim() == 3:
+            lbl = lbl.unsqueeze(1)
+        test_wo.append(wo)
+        test_w.append(w)
+        test_lbl.append(lbl)
+    test_wo = torch.cat(test_wo, 0)
+    test_w = torch.cat(test_w, 0)
+    test_lbl = torch.cat(test_lbl, 0)
+    print(f"Testing samples: {test_w.shape[0]}")
+
+    train_data = torch.utils.data.TensorDataset(train_wo, train_w, train_lbl)
+    test_data = torch.utils.data.TensorDataset(test_wo, test_w, test_lbl)
+    train_loader = torch.utils.data.DataLoader(
+        train_data, batch_size=BATCH_SIZE, shuffle=True,
+        num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True)
+    test_loader = torch.utils.data.DataLoader(
+        test_data, batch_size=BATCH_SIZE, shuffle=False,
+        num_workers=NUM_WORKERS, pin_memory=True, persistent_workers=True)
+
+    timestamp = time.strftime("%Y-%m-%d-%H-%M-%S")
+    save_dir = os.path.join("./runs", f"multiseed_baseline_{timestamp}")
+    os.makedirs(save_dir, exist_ok=True)
+
+    all_results = []
+    for seed in SEEDS:
+        result = train_one_seed(seed, train_loader, test_loader, save_dir)
+        all_results.append(result)
+
+    medians = [r["median"] for r in all_results]
+    p90s = [r["p90"] for r in all_results]
+    p99s = [r["p99"] for r in all_results]
+
+    print(f"\n\n{SEP2}")
+    print(f"  MULTI-SEED RESULTS: BASELINE DLoc")
+    print(SEP2)
+    for i, seed in enumerate(SEEDS):
+        r = all_results[i]
+        print(f"  Seed {seed:<5}: median={r['median']:.4f}m  P90={r['p90']:.3f}m  P99={r['p99']:.3f}m")
+    print(SEP3)
+    print(f"  Median error:  {np.mean(medians):.4f} +/- {np.std(medians):.4f} m")
+    print(f"  P90 error:     {np.mean(p90s):.3f} +/- {np.std(p90s):.3f} m")
+    print(f"  P99 error:     {np.mean(p99s):.3f} +/- {np.std(p99s):.3f} m")
+    print(SEP2)
+
+    report_path = os.path.join(save_dir, "report.txt")
+    with open(report_path, "w") as f:
+        f.write(f"Student: baseline\nMode: standalone\nSeeds: {SEEDS}\nEpochs: {EPOCHS}\n\n")
+        for i, seed in enumerate(SEEDS):
+            r = all_results[i]
+            f.write(f"Seed {seed}: median={r['median']:.4f}m  P90={r['p90']:.3f}m  P99={r['p99']:.3f}m\n")
+        f.write(f"\nMedian: {np.mean(medians):.4f} +/- {np.std(medians):.4f} m\n")
+        f.write(f"P90:    {np.mean(p90s):.3f} +/- {np.std(p90s):.3f} m\n")
+        f.write(f"P99:    {np.mean(p99s):.3f} +/- {np.std(p99s):.3f} m\n")
+    print(f"\nReport saved to: {report_path}")
+
+
+if __name__ == "__main__":
+    main()