From 262c6d9ef3f0461218c1a5e632fc7dd6357e0a20 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 21 Jul 2026 08:48:04 -0500 Subject: feat: add parameter-matched EP comparison --- experiments/run.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) (limited to 'experiments/run.py') diff --git a/experiments/run.py b/experiments/run.py index eea3f91..09e5a30 100644 --- a/experiments/run.py +++ b/experiments/run.py @@ -76,7 +76,8 @@ def train(args): device = args.device torch.manual_seed(args.seed) train_loader, test_loader, n_in, n_out = get_dataset( - args.dataset, batch_size=args.batch_size, device=device) + args.dataset, batch_size=args.batch_size, device=device, + train_limit=args.train_examples or None) args.n_in = n_in net, cfg = build(args, device) @@ -164,6 +165,8 @@ def get_args(): p.add_argument("--residual", type=int, default=0) # skip connections (deep no-BN) p.add_argument("--epochs", type=int, default=15) p.add_argument("--batch_size", type=int, default=128) + p.add_argument("--train_examples", type=int, default=0, + help="0 uses the full training split") p.add_argument("--eta", type=float, default=0.05) p.add_argument("--eta_A", type=float, default=0.02) p.add_argument("--eta_P", type=float, default=0.002) -- cgit v1.2.3