File size: 962 Bytes
f15d29e | 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 | """Train MatterGen with the OneScience materials stack."""
import os
from datetime import datetime
from pathlib import Path
import hydra
import omegaconf
import torch
from omegaconf import OmegaConf
from model.diffusion.config import Config
from model.diffusion.run import main as run_training
EXAMPLE_DIR = Path(__file__).resolve().parent
os.environ.setdefault(
"OUTPUT_DIR",
str(EXAMPLE_DIR / "outputs" / "train" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S")),
)
@hydra.main(
config_path=str(EXAMPLE_DIR / "conf"),
config_name="default",
version_base="1.1",
)
def main(cfg: omegaconf.DictConfig) -> None:
"""Resolve the Hydra config and run MatterGen training."""
torch.set_float32_matmul_precision("high")
config = OmegaConf.merge(OmegaConf.structured(Config), cfg)
OmegaConf.set_readonly(config, True)
print(OmegaConf.to_yaml(cfg, resolve=True))
run_training(config)
if __name__ == "__main__":
main()
|