From b7f78b395a679cfda59fdb0ff61e9d00eff786e0 Mon Sep 17 00:00:00 2001 From: mushroomfire <934313174@qq.com> Date: Thu, 6 Aug 2026 19:55:58 +0300 Subject: [PATCH] Add export_valid_split: write the valid_ratio split as GPUMD-ready files export_valid_split(data_file, valid_ratio, run_seed, output_dir) writes train.xyz + test.xyz reproducing exactly the split that train_nep(valid_ratio=..., run_seed=...) uses internally, with frames copied verbatim (raw text, untouched fields/precision) in input-file order - so the same data partition can be trained in GPUMD and the loss curves compared directly. The split draw is factored into data.valid_split_indices and both trainers now call it, so the trainers and the export can never disagree. Test: reference energies in energy_test.out from a valid_ratio run equal the exported test.xyz energies row for row; partition is verbatim and deterministic. Also restores train_sharded's 'from .data import' line that a refactor script had accidentally dropped (unreleased; DDP smoke test re-run and passing). --- README.md | 6 ++++ releaseNotes.md | 7 ++++ tests/README.md | 2 +- tests/test_run_seed_and_valid.py | 48 ++++++++++++++++++++++++++ torchnep/__init__.py | 2 ++ torchnep/data.py | 58 ++++++++++++++++++++++++++++++++ torchnep/train.py | 24 +++++-------- torchnep/train_sharded.py | 23 ++++--------- 8 files changed, 137 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 3075c89..a5df7b2 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,12 @@ train_nep("nep.in", "train.xyz", output_dir="output") # with a validation set (either form): train_nep("nep.in", "train.xyz", output_dir="output", valid_file="valid.xyz") train_nep("nep.in", "train.xyz", output_dir="output", valid_ratio=0.1) + +# export the exact valid_ratio split as GPUMD-ready files — train the same +# partition in GPUMD (or anything else) and compare loss curves directly: +from torchnep import export_valid_split +export_valid_split("train.xyz", valid_ratio=0.1, run_seed=42, + output_dir="split") # writes split/train.xyz + split/test.xyz ``` ```bash diff --git a/releaseNotes.md b/releaseNotes.md index 7510f72..5e7fedb 100644 --- a/releaseNotes.md +++ b/releaseNotes.md @@ -1,5 +1,12 @@ # Release Notes +## Unreleased + +- **`export_valid_split`**: write the exact `valid_ratio` split + `train_nep` uses as verbatim GPUMD-ready `train.xyz` / `test.xyz` + files, so the same data partition can be trained in GPUMD and the loss + curves compared directly. + ## 1.0.2a1 - **Streaming-only data path**: the preloaded GPU data store and the diff --git a/tests/README.md b/tests/README.md index d2e26b5..6c27ea8 100644 --- a/tests/README.md +++ b/tests/README.md @@ -16,7 +16,7 @@ TEST_DEVICE=cpu pytest tests/ # restrict device (default: cpu + cuda if pre | `test_parsing.py` | Legacy and current `l_max` nep.in / nep.txt parsing. | | `test_ase_calculator.py` | Optional ASE calculator (energy/forces/stress, ZBL split). | | `test_b1_and_gpumd_qscaler.py` | Analytical `b1` offset (residual → 0, `nep_best` ≤ `nep_final`); `use_gpumd_qscaler` reproduces GPUMD's `c=1` q_scaler; `gpumd_init_parameters` re-inits coeffs **and** NN weights uniform(−1,1); L2 (`lambda_2`) shrinks the weights. | -| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch), is off by default, is per-stage (a stage-1 plateau jumps into Stage 2, surviving resume). | +| `test_run_seed_and_valid.py` | `run_seed` reproducibility; `valid_file` / `valid_ratio` deterministic split, best-model selection on validation loss, `*_test.out`, split preserved across resume; `early_stop` fires on a plateau (validation-loss branch), is off by default, is per-stage (a stage-1 plateau jumps into Stage 2, surviving resume). `export_valid_split` reproduces the internal split (verbatim frames, matches `energy_test.out`). | | `test_stream_mode.py` | `StreamDataStore` (the training data store): `collate` is bit-exact vs an independently assembled reference (concatenation + offsets + basis straight from the ops functions); metadata/mask consistency incl. missing channels; prefetched `iter_collated` matches direct collate; 2-rank DDP same-seed reproducibility. The DDP case is local-only: set `TORCHNEP_TEST_DDP=1` (skipped in CI — multi-process rendezvous is unreliable on shared runners). | | `test_compiled_autograd.py` | `CompiledAutogradForce` (make_fx-materialized autograd forces): outputs and parameter gradients (second-order path through the force loss) match eager autograd across batch shapes on one dynamic graph; energy-only calls fall back to eager. CUDA-only — auto-skipped on CPU hosts/CI. | diff --git a/tests/test_run_seed_and_valid.py b/tests/test_run_seed_and_valid.py index 06f53e8..1feacf9 100644 --- a/tests/test_run_seed_and_valid.py +++ b/tests/test_run_seed_and_valid.py @@ -358,3 +358,51 @@ def test_resume_preserves_valid_split(tmp_path): expected = np.array([frames[i]["energy"] / len(frames[i]["species"]) for i in val_idx]) assert np.allclose(ref_resumed, expected, atol=1e-8) + + +def test_export_valid_split_matches_train_nep(tmp_path): + """export_valid_split writes GPUMD-ready train/test files that + reproduce exactly the valid_ratio split train_nep uses: the reference + energies in energy_test.out equal the exported test.xyz energies (in + order), the two files partition the input verbatim, and the export is + deterministic.""" + from torchnep import export_valid_split + nepin, xyz = _write_run_files(tmp_path) + + out = tmp_path / "run" + train_nep(config_file=nepin, data_file=xyz, output_dir=str(out), + device="cpu", precision="float64", print_interval=100, + restart=False, checkpoint_interval=1000, + prediction_interval=1000, run_seed=11, valid_ratio=0.25) + + tp, vp, n_tr, n_va = export_valid_split( + xyz, valid_ratio=0.25, run_seed=11, + output_dir=str(tmp_path / "split")) + + orig = read_xyz(xyz) + tr, va = read_xyz(tp), read_xyz(vp) + assert n_tr == len(tr) and n_va == len(va) + assert len(tr) + len(va) == len(orig) + + # Reference per-atom energies in energy_test.out (col 2) must equal the + # exported test.xyz frames, same order. + rows = np.loadtxt(out / "energy_test.out", ndmin=2) + assert rows.shape[0] == len(va) + ref = rows[:, 1] + exp = np.array([f["energy"] / f["natoms"] for f in va]) + assert np.allclose(ref, exp, atol=1e-10) + + # Verbatim: re-reading the two exports and the original gives the same + # frame multiset (match on energy + natoms fingerprints). + def fp(frames): + return sorted((f["natoms"], round(float(f["energy"]), 10)) + for f in frames) + assert fp(tr) + fp(va) != [] # sanity + assert sorted(fp(tr) + fp(va)) == fp(orig) + + # Deterministic: exporting again is byte-identical. + tp2, vp2, _, _ = export_valid_split( + xyz, valid_ratio=0.25, run_seed=11, + output_dir=str(tmp_path / "split2")) + assert open(tp).read() == open(tp2).read() + assert open(vp).read() == open(vp2).read() diff --git a/torchnep/__init__.py b/torchnep/__init__.py index e232ad7..f51f1b7 100644 --- a/torchnep/__init__.py +++ b/torchnep/__init__.py @@ -16,7 +16,9 @@ from .predict import predict_dataset from .train import train_nep from .train_sharded import train_nep_sharded +from .data import export_valid_split __all__ = [ "predict_dataset", "train_nep", "train_nep_sharded", + "export_valid_split", ] diff --git a/torchnep/data.py b/torchnep/data.py index db128f3..1490f2c 100644 --- a/torchnep/data.py +++ b/torchnep/data.py @@ -462,3 +462,61 @@ def build_neighbor_list_np(positions, cell, cutoff): return (np.concatenate(all_i).astype(np.int64), np.concatenate(all_j).astype(np.int64), np.concatenate(all_rij)) + + +def valid_split_indices(n_frames: int, valid_ratio: float, run_seed: int): + """Train/validation split indices — the exact split ``train_nep`` makes. + + ``train_nep(valid_ratio=r, run_seed=s)`` holds out + ``max(1, round(r * n))`` frames drawn from a dedicated torch generator + seeded with ``run_seed``. This helper is that draw, factored out so the + trainers and :func:`export_valid_split` can never disagree. + + Returns ``(train_idx, valid_idx)`` — both sorted in input-file order. + """ + import torch + if run_seed is None: + raise ValueError("run_seed is required: the split is drawn from it " + "(train_nep uses the same seed to reproduce it)") + if not 0.0 < valid_ratio < 1.0: + raise ValueError(f"valid_ratio must be in (0, 1), got {valid_ratio}") + g = torch.Generator() + g.manual_seed(run_seed) + perm = torch.randperm(n_frames, generator=g).tolist() + n_val = max(1, int(round(valid_ratio * n_frames))) + if n_val >= n_frames: + raise ValueError(f"valid_ratio={valid_ratio} leaves no training " + f"frames ({n_frames} total)") + val_set = set(perm[:n_val]) + train_idx = [i for i in range(n_frames) if i not in val_set] + return train_idx, sorted(val_set) + + +def export_valid_split(data_file: str, valid_ratio: float, run_seed: int, + output_dir: str = "split"): + """Write GPUMD-ready ``train.xyz`` / ``test.xyz`` with train_nep's split. + + Reproduces exactly the validation split that + ``train_nep(data_file, valid_ratio=r, run_seed=s)`` uses internally, so + the exported pair can train the SAME data partition in GPUMD (or any + other code) and loss curves stay comparable. Frames are copied verbatim + (raw text, untouched fields and precision), in input-file order. + + Returns ``(train_path, test_path, n_train, n_valid)``. + """ + import os + with open(data_file) as f: + blocks = _split_frames(f.readlines()) + train_idx, val_idx = valid_split_indices(len(blocks), valid_ratio, + run_seed) + os.makedirs(output_dir, exist_ok=True) + train_path = os.path.join(output_dir, "train.xyz") + test_path = os.path.join(output_dir, "test.xyz") + src = os.path.abspath(data_file) + for path, idxs in ((train_path, train_idx), (test_path, val_idx)): + if os.path.abspath(path) == src: + raise ValueError(f"output would overwrite the input: {src}") + with open(path, "w") as out: + for k in idxs: + out.writelines(blocks[k]) + return train_path, test_path, len(train_idx), len(val_idx) diff --git a/torchnep/train.py b/torchnep/train.py index 781c4eb..f974dc5 100644 --- a/torchnep/train.py +++ b/torchnep/train.py @@ -41,7 +41,7 @@ from torch.optim.swa_utils import AveragedModel from .model import NEPModel, slim_model, gpumd_init_parameters -from .data import read_xyz, parse_nep_in, build_neighbor_list_np +from .data import read_xyz, parse_nep_in, valid_split_indices, build_neighbor_list_np from . import ops from . import __version__ from .predict import predict_from_store @@ -1324,22 +1324,14 @@ def _log(msg=""): _log(f" read {len(valid_frames)} validation structures " f"from {valid_file}") elif valid_ratio is not None: - if not 0.0 < valid_ratio < 1.0: - raise ValueError(f"valid_ratio must be in (0, 1), " - f"got {valid_ratio}") - g = torch.Generator() - g.manual_seed(run_seed) - perm = torch.randperm(len(frames), generator=g).tolist() - n_val = max(1, int(round(valid_ratio * len(frames)))) - if n_val >= len(frames): - raise ValueError(f"valid_ratio={valid_ratio} leaves no " - f"training frames ({len(frames)} total)") - val_idx = sorted(perm[:n_val]) - val_set = set(val_idx) + # The same draw export_valid_split reproduces (see data.py) — the + # split can be exported as GPUMD-ready train.xyz/test.xyz files. + train_idx, val_idx = valid_split_indices(len(frames), valid_ratio, + run_seed) valid_frames = [frames[i] for i in val_idx] - frames = [f for i, f in enumerate(frames) if i not in val_set] - _log(f" valid_ratio={valid_ratio}: held out {n_val} frames for " - f"validation, {len(frames)} remain for training " + frames = [frames[i] for i in train_idx] + _log(f" valid_ratio={valid_ratio}: held out {len(val_idx)} frames " + f"for validation, {len(frames)} remain for training " f"(split drawn from run_seed)") # Single-GPU: per-epoch shuffle is done at iteration time via diff --git a/torchnep/train_sharded.py b/torchnep/train_sharded.py index 08d48aa..393cdab 100644 --- a/torchnep/train_sharded.py +++ b/torchnep/train_sharded.py @@ -42,7 +42,7 @@ import torch.nn as nn from .model import NEPModel, gpumd_init_parameters -from .data import read_xyz, parse_nep_in +from .data import read_xyz, parse_nep_in, valid_split_indices from . import ops from . import __version__ from .predict import predict_from_store_sharded @@ -423,22 +423,13 @@ def _log(msg=""): _log(f" read {len(valid_frames)} validation structures " f"from {valid_file}") elif valid_ratio is not None: - if not 0.0 < valid_ratio < 1.0: - raise ValueError(f"valid_ratio must be in (0, 1), " - f"got {valid_ratio}") - vg = torch.Generator() - vg.manual_seed(run_seed) - vperm = torch.randperm(len(frames), generator=vg).tolist() - n_val = max(1, int(round(valid_ratio * len(frames)))) - if n_val >= len(frames): - raise ValueError(f"valid_ratio={valid_ratio} leaves no " - f"training frames ({len(frames)} total)") - val_idx = sorted(vperm[:n_val]) - val_set = set(val_idx) + # Same draw as export_valid_split / train_nep (see data.py). + train_idx, val_idx = valid_split_indices(len(frames), valid_ratio, + run_seed) valid_frames = [frames[i] for i in val_idx] - frames = [f for i, f in enumerate(frames) if i not in val_set] - _log(f" valid_ratio={valid_ratio}: held out {n_val} frames for " - f"validation, {len(frames)} remain for training " + frames = [frames[i] for i in train_idx] + _log(f" valid_ratio={valid_ratio}: held out {len(val_idx)} frames " + f"for validation, {len(frames)} remain for training " f"(split drawn from run_seed)") n_total = len(frames)