mschoi/ai-rfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use final 100-run RMSD as training outcome.

Browse Source 
Switch success tracking from train loss to final mean RMSD over 100 rollout predictions, and persist latest/best metric artifacts for gated train.py commits.

Made-with: Cursor
This commit is contained in:
demian3b
2026-04-16 16:57:52 +09:00
parent 7a92652289
commit 49baf72a4a
4 changed files with 142 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
BEST_PRACTICE.json
View File

@@ -1,8 +1,8 @@
{
"best_mean_rmsd_100": 999999.0,
"best_mean_rmsd_100": 2.5936938130855562,
"num_runs": 100,
"timestamp_utc": "1970-01-01T00:00:00Z",
"command": "",
"notes": "Bootstrap placeholder. Replace after first measured run.",
"timestamp_utc": "2026-04-16T07:56:51.073448+00:00",
"command": "train.py --sdf /data/demian_dev/toy/sample.sdf --epochs 10 --batch-size 64 --num-workers 8 --prefetch-factor 8 --hidden 256 --gcn-layers 6 --accum 2 --eval-runs 100 --rollout-steps 100 --out-dir /tmp/ai_rfm_eval",
"notes": "Synced to latest final test result.",
"updated_by_commit": ""
}
}

5
README.md
View File

@@ -4,7 +4,7 @@ RFM overfitting sandbox for a single ligand sample, with hard quality gates.
## Environment first (UV, cu126 only)
1. Ensure Python 3.12 is available.
1. Ensure Python 3.10 is available.
2. Install env and deps:
- `uv sync`
3. Install git hooks:
@@ -18,7 +18,7 @@ This repository is intentionally pinned to CUDA 12.6 PyTorch wheels and matching
- Commits touching `train.py` must include:
- `reports/latest_eval.json`
- `BEST_PRACTICE.json`
- better or equal `mean_rmsd_100` compared to previous best (enforced by pre-commit).
- strictly better `mean_rmsd_100` compared to previous best (enforced by pre-commit).
## Evaluation target
@@ -36,3 +36,4 @@ This repository is intentionally pinned to CUDA 12.6 PyTorch wheels and matching
## Attempt Log
- 2026-04-16: Bootstrapped docs/environment policy and cu126 UV config. Added best-practice/performance gating scaffolding before the next training run.
- 2026-04-16: Updated `train.py` to use final test metric as source of truth (`mean_rmsd_100` from 100 rollout predictions) and removed train-loss based best checkpoint tracking. Current measured `mean_rmsd_100=2.593694`.

10
reports/latest_eval.json
View File

@@ -1,7 +1,7 @@
{
"mean_rmsd_100": 999999.0,
"mean_rmsd_100": 2.5936938130855562,
"num_runs": 100,
"timestamp_utc": "1970-01-01T00:00:00Z",
"command": "",
"notes": "Bootstrap placeholder. Replace with real measured evaluation output."
}
"timestamp_utc": "2026-04-16T07:56:51.073448+00:00",
"command": "train.py --sdf /data/demian_dev/toy/sample.sdf --epochs 10 --batch-size 64 --num-workers 8 --prefetch-factor 8 --hidden 256 --gcn-layers 6 --accum 2 --eval-runs 100 --rollout-steps 100 --out-dir /tmp/ai_rfm_eval",
"notes": "Final test metric from 100 rollout predictions."
}

150
train.py
View File

@@ -7,6 +7,7 @@ Consolidated RFM toy training + visualization from sandbox_rfm.ipynb.
from __future__ import annotations
import argparse
import json
import os
import pickle
import sys
@@ -14,6 +15,7 @@ import time
from collections import deque
from copy import deepcopy
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any
import torch
@@ -305,6 +307,80 @@ def save_video_visualization(
writer.close()
# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------
def kabsch_rmsd(pred: torch.Tensor, target: torch.Tensor) -> float:
"""RMSD after rigid alignment (Kabsch), shape: [N, 3]."""
p = pred - pred.mean(dim=0, keepdim=True)
q = target - target.mean(dim=0, keepdim=True)
h = p.T @ q
u, _, v_t = torch.linalg.svd(h)
r = v_t.T @ u.T
if torch.det(r) < 0:
v_t[-1, :] *= -1
r = v_t.T @ u.T
p_aligned = p @ r
return float(torch.sqrt(torch.mean(torch.sum((p_aligned - q) ** 2, dim=-1))).item())
def graph_from_state_cpu(sample: Data, center: torch.Tensor, rot: torch.Tensor, torsion: torch.Tensor) -> Data:
g = deepcopy(sample)
g.coords = sample.gt_coords.clone()
apply_trans(g, center)
apply_rot_matrix(g, rot)
apply_torsion(g, torsion, g.movable_mask, g.bond_endpoints)
return g
def evaluate_mean_rmsd_100(
model: nn.Module,
sample: Data,
target_coords: torch.Tensor,
device: torch.device,
num_runs: int = 100,
rollout_steps: int = 100,
) -> float:
"""Run 100 rollouts and average final-structure RMSD."""
model.eval()
nb = sample.bond_endpoints.shape[0]
centers = []
rots = []
torsions = []
for _ in range(num_runs):
centers.append(generate_random_trans(torch.zeros(3), torch.tensor([3.0, 3.0, 3.0])))
rots.append(generate_rotation_matrix(generate_random_quaternion()))
torsions.append(generate_random_torsion_angles(nb))
center_t = torch.stack(centers, dim=0) # [R,3]
rot_t = torch.stack(rots, dim=0) # [R,3,3]
torsion_t = torch.stack(torsions, dim=0) # [R,B]
dt = 1.0 / max(1, rollout_steps - 1)
for k in range(rollout_steps):
t_val = k / max(1, rollout_steps - 1)
graphs = [
graph_from_state_cpu(sample, center_t[i], rot_t[i], torsion_t[i])
for i in range(num_runs)
]
batch = Batch.from_data_list(graphs).to(device)
tb = torch.full((num_runs, 1), t_val, device=device, dtype=torch.float32)
with torch.no_grad():
v = model(batch, tb)
vc = v.center.detach().cpu()
vw = v.rot_omega.detach().cpu()
vt = v.torsion.detach().cpu()
center_t = center_t + vc * dt
torsion_t = torsion_t + vt * dt
for i in range(num_runs):
rot_t[i] = rot_t[i] @ so3_exp(vw[i] * dt)
rmsds = []
for i in range(num_runs):
final_graph = graph_from_state_cpu(sample, center_t[i], rot_t[i], torsion_t[i])
rmsds.append(kabsch_rmsd(final_graph.coords, target_coords))
return float(sum(rmsds) / len(rmsds))
# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
@@ -535,6 +611,18 @@ def main() -> int:
default=1_000_000,
help="Synthetic Dataset length (only affects shuffle / epoch size)",
)
parser.add_argument(
"--eval-runs",
type=int,
default=100,
help="Number of rollout tests for final RMSD report",
)
parser.add_argument(
"--rollout-steps",
type=int,
default=100,
help="Integration steps per rollout during final test",
)
args = parser.parse_args()
torch.manual_seed(args.seed)
@@ -575,10 +663,7 @@ def main() -> int:
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("device:", device)
print("model param device:", next(model.parameters()).device)
print("cuda available:", torch.cuda.is_available())
print(f"model: hidden={args.hidden} layers={args.gcn_layers} params~={n_params:,} accum={args.accum}")
print(f"device={device} hidden={args.hidden} layers={args.gcn_layers} params={n_params:,}")
nw = args.num_workers
if nw < 0:
@@ -615,8 +700,6 @@ def main() -> int:
f"persistent_workers={persistent}"
)
best_loss_gpu = torch.tensor(float("inf"), device=device)
best_state: dict[str, torch.Tensor] | None = None
t0 = time.perf_counter()
log_every = max(1, args.epochs // 10)
accum = max(1, args.accum)
@@ -632,9 +715,6 @@ def main() -> int:
tw = tw.to(device=device, dtype=torch.float32, non_blocking=True)
tt = tt.to(device=device, dtype=torch.float32, non_blocking=True)
if epoch == 0:
assert batched.coords.is_cuda == (device.type == "cuda"), "batched coords should match cuda flag"
pred = model(batched, time_b)
mse = (
F.mse_loss(pred.center, tc)
@@ -647,25 +727,53 @@ def main() -> int:
optimizer.step()
avg_mse = sum_mse / accum
with torch.no_grad():
if avg_mse < best_loss_gpu:
best_loss_gpu = avg_mse.clone()
m = model._orig_mod if hasattr(model, "_orig_mod") else model
best_state = {k: v.detach().cpu().clone() for k, v in m.state_dict().items()}
if epoch % log_every == 0 or epoch == args.epochs - 1:
loss_show = float(avg_mse.detach().cpu())
best_show = float(best_loss_gpu.detach().cpu())
print(f"epoch {epoch:6d} loss {loss_show:.6f} best {best_show:.6f}")
print(f"epoch {epoch:6d} train_mse {loss_show:.6f}")
if device.type == "cuda":
torch.cuda.synchronize()
best_loss = float(best_loss_gpu.detach().cpu())
print(f"done in {time.perf_counter() - t0:.1f}s best_loss={best_loss:.6f}")
print(f"training done in {time.perf_counter() - t0:.1f}s")
if best_state is not None:
m_load = model._orig_mod if hasattr(model, "_orig_mod") else model
m_load.load_state_dict(best_state)
# Final test metric: 100 rollouts mean final-structure RMSD
final_mean_rmsd_100 = evaluate_mean_rmsd_100(
model,
sample,
sample.gt_coords,
device=device,
num_runs=args.eval_runs,
rollout_steps=args.rollout_steps,
)
print(f"final_test_mean_rmsd_{args.eval_runs}: {final_mean_rmsd_100:.6f}")
reports_dir = os.path.join(here, "reports")
os.makedirs(reports_dir, exist_ok=True)
best_path = os.path.join(here, "BEST_PRACTICE.json")
latest_path = os.path.join(reports_dir, "latest_eval.json")
timestamp = datetime.now(timezone.utc).isoformat()
cmd = " ".join(sys.argv)
latest_report = {
"mean_rmsd_100": float(final_mean_rmsd_100),
"num_runs": int(args.eval_runs),
"timestamp_utc": timestamp,
"command": cmd,
"notes": "Final test metric from 100 rollout predictions.",
}
with open(latest_path, "w", encoding="utf-8") as f:
json.dump(latest_report, f, indent=2)
best_report = {
"best_mean_rmsd_100": float(final_mean_rmsd_100),
"num_runs": int(args.eval_runs),
"timestamp_utc": timestamp,
"command": cmd,
"notes": "Synced to latest final test result.",
"updated_by_commit": "",
}
with open(best_path, "w", encoding="utf-8") as f:
json.dump(best_report, f, indent=2)
# --- Visualizations ---
sample.coords = sample.gt_coords.clone()