diff options
Diffstat (limited to 'experiments/diagnose_kp_traffic_nonfinite.py')
| -rw-r--r-- | experiments/diagnose_kp_traffic_nonfinite.py | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/experiments/diagnose_kp_traffic_nonfinite.py b/experiments/diagnose_kp_traffic_nonfinite.py index 99a4e99..7b7389f 100644 --- a/experiments/diagnose_kp_traffic_nonfinite.py +++ b/experiments/diagnose_kp_traffic_nonfinite.py @@ -90,6 +90,7 @@ def main(): choices=("nlms20", "closed_form"), default="nlms20") parser.add_argument("--predictor_every", type=int, default=16) + parser.add_argument("--stability_margin", type=float, default=0.0) parser.add_argument("--device", default="cuda") parser.add_argument("--max_steps", type=int, default=352) parser.add_argument("--out", required=True) @@ -98,6 +99,8 @@ def main(): raise ValueError("max_steps must be positive") if args.predictor_every < 0: raise ValueError("predictor_every must be nonnegative") + if args.stability_margin < 0: + raise ValueError("stability_margin must be nonnegative") torch.manual_seed(0) if str(args.device).startswith("cuda"): @@ -137,7 +140,8 @@ def main(): warmup_mse = [] if args.predictor_mode == "closed_form": closed_form_fit = net.predictor_closed_form_fit( - calibration_forward["hiddens"]) + calibration_forward["hiddens"], + stability_margin=args.stability_margin) warmup_mse.append(closed_form_fit["mse"]) else: iterator = iter(train) @@ -228,6 +232,7 @@ def main(): "rule": args.rule, "predictor_mode": args.predictor_mode, "predictor_every": args.predictor_every, + "stability_margin": args.stability_margin, "max_steps": args.max_steps, "split": split, "traffic_calibration": calibration, |