feat: add parameter-matched EP comparison

1 files changed, 4 insertions, 1 deletions

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)