diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 74d6a3b..b85ac96 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -211,6 +211,7 @@ def ep_step(x, y): drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( sum(float(b.norm()) for b in zs_free), 1e-9) if (not math.isfinite(drift)) or (drift > 0.5 and not args.noguard): + GOV['skd'] = GOV.get('skd', 0) + 1 # drift-guard reject (relaxation non-convergence) for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False GOV['drift'] = drift @@ -226,6 +227,7 @@ def ep_step(x, y): GOV['gema'] = 0.99 * GOV['gema'] + 0.01 * gn # EMA always updates (frozen-ref bugfix) GOV['gn'] = gn if not math.isfinite(gn) or (gn > 8 * GOV['gema'] and not args.noguard): + GOV['skg'] = GOV.get('skg', 0) + 1 # gn-EMA-guard reject (gradient-magnitude spike) for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False for p, g in zip(all_params, gs): @@ -287,7 +289,7 @@ for step in range(start_step, args.steps + 1): val = evaluate(); best = min(best, val) gtag = '' if math.isnan(gcos) else f' cos={gcos:.4f}' print(f'step {step:5d}/{args.steps} | train {ce:.4f} val {val:.4f} (best {best:.4f}) ' - f'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} ' + f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}){gtag} ' f'drift={GOV["drift"]:.3f} gn={GOV["gn"]:.2e} sig={GOV["sig"]:.1f} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True) if wb is not None: try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best, 'beta_t': beta_t, |