File size: 7,405 Bytes
d7e0e08 | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """
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()
|