File size: 2,739 Bytes
096347b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | #!/usr/bin/env python3
"""Fit OC20 energy element references before starting model training."""
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import yaml
from onescience.utils.uma.normalization.element_references import (
fit_linear_references,
)
from train import _expand_path, _loader
def fit_references(
config_path: str | Path,
output_path: str | Path,
batch_size: int = 32,
num_batches: int | None = None,
) -> Path:
"""Fit and save legacy ``coeff`` references consumed by the OC20 YAML."""
if batch_size < 2:
raise ValueError("reference fitting batch_size must be at least 2")
if num_batches is not None and num_batches < 1:
raise ValueError("num_batches must be positive")
config_path = Path(config_path).expanduser().resolve()
with config_path.open(encoding="utf-8") as handle:
config = yaml.safe_load(handle) or {}
train_path = _expand_path(config.get("train"))
if not train_path:
raise ValueError(f"training data is not configured in {config_path}")
loader = _loader(
train_path,
batch_size=batch_size,
workers=int(config.get("workers", 0)),
max_samples=config.get("max_train_samples"),
train=False,
seed=int(config.get("seed", 0)),
max_atoms=config.get("max_atoms"),
load_balancing=False,
)
references = fit_linear_references(
targets=["energy"],
dataset=loader.dataset,
batch_size=batch_size,
num_batches=num_batches,
num_workers=int(config.get("workers", 0)),
shuffle=False,
)
output_path = Path(output_path).expanduser().resolve()
if output_path.suffix != ".npz":
raise ValueError("output path must use the .npz extension")
output_path.parent.mkdir(parents=True, exist_ok=True)
coefficients = (
references["energy"].element_references.detach().cpu().numpy()
)
np.savez(output_path, coeff=coefficients)
print(f"Saved energy element references to {output_path}")
return output_path
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", required=True, help="OC20 training YAML")
parser.add_argument("--output", required=True, help="output .npz path")
parser.add_argument("--batch-size", type=int, default=32)
parser.add_argument("--num-batches", type=int)
return parser.parse_args()
def main() -> None:
args = _parse_args()
fit_references(
args.config,
args.output,
batch_size=args.batch_size,
num_batches=args.num_batches,
)
if __name__ == "__main__":
main()
|