71 lines
2.9 KiB
Python
71 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
# ---------------------------------------------------------------
|
|
# Fixed-orthogonal-mask scheme across user counts U = 1..6, for
|
|
# overlaying the independent reference on Fig. 2 and Fig. 3 of the
|
|
# WCL revision.
|
|
#
|
|
# Reuses the EXACT training/eval path of eval_task_oriented.py
|
|
# (fixed QR masks, transceiver trained with MSE+CosSim, seed 0,
|
|
# matched 100-epoch x 200-step budget) so the U=4 numbers reproduce
|
|
# Table II. Evaluates the full SNR sweep per U so the same file
|
|
# serves Fig 3(a) (CosSim vs SNR, U=4), Fig 2(b) (U=4, 20 dB) and
|
|
# Fig 3(b) (throughput vs U at 10 dB).
|
|
#
|
|
# Output: ../results_sweeps/task_oriented/fixed_orth_byU.csv
|
|
# ---------------------------------------------------------------
|
|
import os, csv, argparse
|
|
from types import SimpleNamespace
|
|
import torch
|
|
|
|
from eval_task_oriented import (load_emb, train_transceiver, sweep_metrics,
|
|
build_fixed_orthogonal_masks)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--users-list", type=int, nargs="+",
|
|
default=[1, 2, 3, 4, 5, 6])
|
|
ap.add_argument("--epochs", type=int, default=100)
|
|
ap.add_argument("--steps-per-epoch", type=int, default=200)
|
|
ap.add_argument("--eval-trials", type=int, default=200)
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--embed-file", type=str, default="../bert_agnews_8000.pt")
|
|
ap.add_argument("--out", type=str,
|
|
default="../results_sweeps/task_oriented/fixed_orth_byU.csv")
|
|
a = ap.parse_args()
|
|
|
|
device = torch.device("mps" if torch.backends.mps.is_available()
|
|
else ("cuda" if torch.cuda.is_available() else "cpu"))
|
|
print(f"[INFO] device={device}")
|
|
d_bert, mux = 768, 4
|
|
d_s = d_bert * mux
|
|
emb = load_emb(a.embed_file, d_bert, 8000, device)
|
|
|
|
rows_all = []
|
|
for U in a.users_list:
|
|
print(f"\n=== Fixed-Orth scheme, U={U} ===", flush=True)
|
|
args = SimpleNamespace(
|
|
users=U, users_max=8, mux_factor=mux, d_bert=d_bert,
|
|
hidden=256, rank=64, channel="rayleigh",
|
|
train_snr=[0, 5, 10, 15, 20, 25], lr=1e-3, ce_tau=16.0,
|
|
epochs=a.epochs, steps_per_epoch=a.steps_per_epoch,
|
|
eval_trials=a.eval_trials, seed=a.seed)
|
|
Mfix = build_fixed_orthogonal_masks(U, d_s, seed=0)
|
|
trx, mfn = train_transceiver(emb, args, device, "fixed",
|
|
fixed_masks=Mfix)
|
|
rows = sweep_metrics(trx, emb, mfn, device, U, a.eval_trials)
|
|
for r in rows:
|
|
rows_all.append((U, *r))
|
|
|
|
os.makedirs(os.path.dirname(a.out), exist_ok=True)
|
|
with open(a.out, "w", newline="") as f:
|
|
w = csv.writer(f)
|
|
w.writerow(["users", "snr_db", "cos_sim", "orthogonality", "top1_acc"])
|
|
for r in rows_all:
|
|
w.writerow(r)
|
|
print(f"\n[DONE] wrote {a.out}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|