diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py index b4de085..00fb7a5 100644 --- a/scripts/train_merge_code.py +++ b/scripts/train_merge_code.py @@ -30,7 +30,6 @@ OUT = Path(os.environ.get("LOOP_OUT", Path(__file__).resolve().parent.parent / "results-loop")) STEPS = 600 BATCH = 4 -LR = 1e-3 WARMUP = 20 MAX_TOK = 512 K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] @@ -51,13 +50,16 @@ ap.add_argument("--deepk", type=int, default=0, help="scale curriculum depths by deepk/4 (e.g. 16 -> 2/8/16)") ap.add_argument("--bptt", type=int, default=0, help="truncated BPTT: grads only through last N iterations") +ap.add_argument("--lr", type=float, default=1e-3) ARGS = ap.parse_args() SEED = ARGS.seed +LR = ARGS.lr SUFFIX = ((f"_s{SEED}" if SEED else "") + (f"_p{ARGS.pause}" if ARGS.pause else "") + (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "") + ("_ad" if ARGS.adaptive else "") - + (f"_dk{ARGS.deepk}" if ARGS.deepk else "")) + + (f"_dk{ARGS.deepk}" if ARGS.deepk else "") + + (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")) PAUSE_ID = 6 #