diff options
Diffstat (limited to 'src/train.py')
| -rw-r--r-- | src/train.py | 20 |
1 files changed, 18 insertions, 2 deletions
diff --git a/src/train.py b/src/train.py index 975598d..46bd8e7 100644 --- a/src/train.py +++ b/src/train.py @@ -1,8 +1,10 @@ """Training loop. Usage: python -m src.train [model] [seed] [--flag ...]""" import csv +import json import math import os import random +import sys import numpy as np import torch @@ -62,8 +64,21 @@ def main() -> None: cfg = parse_args() set_seed(cfg.seed) out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}") + csv_path = os.path.join(out_dir, "metrics.csv") + if os.path.exists(csv_path): + raise SystemExit(f"REFUSING to rerun in place: {csv_path} exists. Use a fresh --out_dir " + f"(reruns would corrupt the CSV and checkpoint provenance).") os.makedirs(out_dir, exist_ok=True) cfg.save(os.path.join(out_dir, "config.json")) + with open(os.path.join(out_dir, "run_meta.json"), "w") as fh: + json.dump({ + "python": sys.version.split()[0], + "torch": torch.__version__, + "numpy": np.__version__, + "device": "cpu", + "torch_threads": torch.get_num_threads(), + "cmd": sys.argv, + }, fh, indent=2) train_in, val_in = get_splits(cfg) train_ex = build_examples(train_in, cfg) @@ -75,9 +90,8 @@ def main() -> None: opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay) ce = nn.CrossEntropyLoss(reduction="none") - log_examples = sorted(val_ex, key=lambda x: (len(str(x)), x))[: cfg.log_n_examples] + log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples] - csv_path = os.path.join(out_dir, "metrics.csv") fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc", "val_em", "mean_halt_steps", "log_examples", "param_count"] best_val_em = -1.0 @@ -152,6 +166,8 @@ def main() -> None: _eval_pass(loss.detach(), halt.detach() if halt is not None else None) if done: break + if step >= cfg.max_train_steps: + break torch.save(model.state_dict(), os.path.join(out_dir, "last.pt")) print(f"DONE steps={step} best_val_em={best_val_em:.4f}") |