diff --git a/experiments/k3/attnres_gradient/train.py b/experiments/k3/attnres_gradient/train.py index f95e7c4..9e85504 100644 --- a/experiments/k3/attnres_gradient/train.py +++ b/experiments/k3/attnres_gradient/train.py @@ -602,6 +602,22 @@ def main() -> None: ) gradient_gate = None diagnostics = [{"step": 0, **initial_diagnostic}] + print( + json.dumps( + { + "event": "diagnostic", + "step": 0, + "architecture": args.architecture, + "depth": args.depth, + "validation_bpc": evaluations[0]["bits_per_byte"], + "activation_gradient_cv": initial_diagnostic[ + "activation_grad_statistics" + ]["population_cv"], + }, + sort_keys=True, + ), + flush=True, + ) model.zero_grad(set_to_none=True) training_history: list[dict[str, float | int]] = [] @@ -664,6 +680,22 @@ def main() -> None: ), } ) + print( + json.dumps( + { + "event": "diagnostic", + "step": step, + "architecture": args.architecture, + "depth": args.depth, + "validation_bpc": evaluations[-1]["bits_per_byte"], + "activation_gradient_cv": diagnostics[-1][ + "activation_grad_statistics" + ]["population_cv"], + }, + sort_keys=True, + ), + flush=True, + ) model.zero_grad(set_to_none=True) model.train()