diff options
Diffstat (limited to 'experiments/run.py')
| -rw-r--r-- | experiments/run.py | 5 |
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) |