From aacb883bc6e6ed7e95ba29a906ab1a68c54fd35f Mon Sep 17 00:00:00 2001 From: Luis Barroso-Luque Date: Sat, 13 Jul 2024 06:18:53 +0300 Subject: [PATCH] [BE] lint/test/docs on PR and push to main (#742) * lint/test on PR and push to main * Run docs CI on every push to main * Fail docs when a notebook doesn't build * skip eval_metrics fix in tutorial_utils * when writing extxyz in docs, specific force as a column explicitly * task.dataset -> dataset.format * small fixes in config in tutorial * make train dataset optional if only used test dataset * more small fixes if you don't specify a train dataset * loss_fns -> loss_functions in tutorial * eval_metrics -> evaluation_metrics in OCP tutorial * FixAtoms as int * don't clear train dataset if present but src is missing * invert train dataset load logic * fix logic * bool mask * require docs test on main * pull_request_review to trigger doc build * use FixAtoms fix from stress-relaxations branch * sneak in a push trigger to test docs * upload only html artifact * sneaky push * fix artifact path * fix job names * no deploy docs on push * add devcontainer settings * small fix to legacy tutorial * small fix to legacy tutorial * another small tutorial fix * add rest of tutorial back * clean up logic in load_datasets * typo * logic typo in load_datasets * convert all md to ipynb on launch * add minimum spec on devcontainer to make sure everything runs * include subdirectories in conversion * 758 fix * minor fixes * relax otf * typo * try removing comments in build_docs yml --------- Co-authored-by: zulissimeta <122578103+zulissimeta@users.noreply.github.com> Co-authored-by: Zack Ulissi Co-authored-by: Muhammed Shuaibi --- .devcontainer/devcontainer.json | 22 +++++ .devcontainer/postCreateCommand.sh | 11 +++ .../workflows/{docs.yml => build_docs.yml} | 25 ++--- .github/workflows/deploy_docs.yml | 34 +++++++ .github/workflows/lint.yml | 5 +- .github/workflows/test.yml | 4 +- docs/_config.yml | 2 + docs/core/fine-tuning/fine-tuning-oxides.md | 4 +- docs/core/inference.md | 8 +- docs/legacy_tutorials/OCP_Tutorial.md | 40 ++++---- docs/legacy_tutorials/data_visualization.md | 2 +- .../advanced/fine-tuning-in-python.md | 4 +- .../core/common/relaxation/ase_utils.py | 2 +- .../relaxation/optimizers/lbfgs_torch.py | 4 +- src/fairchem/core/common/tutorial_utils.py | 4 - src/fairchem/core/trainers/base_trainer.py | 93 +++++++++++-------- src/fairchem/core/trainers/ocp_trainer.py | 4 +- 17 files changed, 180 insertions(+), 88 deletions(-) create mode 100644 .devcontainer/devcontainer.json create mode 100644 .devcontainer/postCreateCommand.sh rename .github/workflows/{docs.yml => build_docs.yml} (59%) create mode 100644 .github/workflows/deploy_docs.yml diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 000000000..77e52140c --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,22 @@ +{ + "image": "mcr.microsoft.com/devcontainers/anaconda:0.205.0-3", + "postCreateCommand": "bash ./.devcontainer/postCreateCommand.sh", + "forwardPorts": [ + 8888 + ], + "postStartCommand": "nohup bash -c 'jupyter notebook --ip=0.0.0.0 --port=8888 --allow-root --NotebookApp.token='' &'", + // Configure tool-specific properties. + "customizations": { + // Configure properties specific to VS Code. + "vscode": { + // Add the IDs of extensions you want installed when the container is created. + "extensions": [ + "ms-python.python", + "ms-python.vscode-pylance" + ] + } + }, + "hostRequirements": { + "memory": "16gb", + } +} \ No newline at end of file diff --git a/.devcontainer/postCreateCommand.sh b/.devcontainer/postCreateCommand.sh new file mode 100644 index 000000000..c73e2c336 --- /dev/null +++ b/.devcontainer/postCreateCommand.sh @@ -0,0 +1,11 @@ +#!/bin/bash +if [ -f packages/requirements.txt ]; then pip install -r packages/requirements.txt; fi +if [ -f packages/requirements-optional.txt ]; then pip install -r packages/requirements-optional.txt; fi +pip install -e packages/fairchem-core[dev] +pip install -e packages/fairchem-data-oc[dev] +pip install -e packages/fairchem-demo-ocpapi[dev] +pip install -e packages/fairchem-applications-cattsunami +pip install jupytext + +# Convert all .md docs to ipynb for easy viewing in vscode later! +find ./docs -name '*.md' -exec jupytext --to ipynb {} \; \ No newline at end of file diff --git a/.github/workflows/docs.yml b/.github/workflows/build_docs.yml similarity index 59% rename from .github/workflows/docs.yml rename to .github/workflows/build_docs.yml index 70ef772b0..276024d2b 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/build_docs.yml @@ -1,12 +1,14 @@ -name: documentation +# build the documentation and upload the built artifact +name: build documentation on: workflow_call: workflow_dispatch: + pull_request_review: + types: [submitted, edited] -# This job installs dependencies, build the book, and pushes it to `gh-pages` jobs: - build-and-deploy: + build: runs-on: ubuntu-latest steps: @@ -36,17 +38,8 @@ jobs: run: | jupyter-book build docs - # Deploy the book's HTML to gh-pages branch # TODO remove once ODAC link updated - - name: Deploy to ghpages branch - uses: peaceiris/actions-gh-pages@v3 + - name: Upload documentation artifact + uses: actions/upload-artifact@v4 with: - github_token: ${{ secrets.GITHUB_TOKEN }} - publish_dir: docs/_build/html - - - name: Deploy to fair-chem.github.io - uses: peaceiris/actions-gh-pages@v4 - with: - deploy_key: ${{ secrets.ACTIONS_DEPLOY_KEY }} - external_repository: FAIR-Chem/fair-chem.github.io - publish_branch: gh-pages - publish_dir: docs/_build/html + name: docs-html + path: docs/_build/html/* diff --git a/.github/workflows/deploy_docs.yml b/.github/workflows/deploy_docs.yml new file mode 100644 index 000000000..c46e12619 --- /dev/null +++ b/.github/workflows/deploy_docs.yml @@ -0,0 +1,34 @@ +name: build and deploy docs + +on: + workflow_call: + workflow_dispatch: + +jobs: + build: + uses: ./.github/workflows/build_docs.yml + + deploy: + needs: build + runs-on: ubuntu-latest + + steps: + - name: Download docs artifact + uses: actions/download-artifact@v4 + with: + name: docs-html + path: docs-html/ + + - name: Deploy to ghpages branch + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: docs-html + + - name: Deploy to fair-chem.github.io + uses: peaceiris/actions-gh-pages@v4 + with: + deploy_key: ${{ secrets.ACTIONS_DEPLOY_KEY }} + external_repository: FAIR-Chem/fair-chem.github.io + publish_branch: gh-pages + publish_dir: docs-html diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 2e5be1d36..45cd6b1cf 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -1,8 +1,11 @@ name: lint on: - push: + workflow_call: workflow_dispatch: + pull_request: + branches: [main] + push: jobs: lint: diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9fa9c1206..d751b29d0 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,9 +2,11 @@ name: tests on: workflow_call: + workflow_dispatch: pull_request: branches: [main] - workflow_dispatch: + push: + branches: [main] jobs: test: diff --git a/docs/_config.yml b/docs/_config.yml index a47c0e116..bd3473e1b 100644 --- a/docs/_config.yml +++ b/docs/_config.yml @@ -52,3 +52,5 @@ sphinx: autosummary_generate: True autoapi_dirs: ['../src/fairchem/core','../src/fairchem/data','../src/fairchem/applications/AdsorbML/adsorbml','../src/fairchem/applications/cattsunami','../src/fairchem/demo/ocpapi'] autoapi_python_use_implicit_namespaces: True + nb_execution_show_tb: True + nb_execution_raise_on_error: True diff --git a/docs/core/fine-tuning/fine-tuning-oxides.md b/docs/core/fine-tuning/fine-tuning-oxides.md index 574989a33..d01b8fcbd 100644 --- a/docs/core/fine-tuning/fine-tuning-oxides.md +++ b/docs/core/fine-tuning/fine-tuning-oxides.md @@ -207,21 +207,23 @@ yml = generate_yml_config(checkpoint_path, 'config.yml', 'optim.loss_force', # the checkpoint setting causes an error 'dataset', 'test_dataset', 'val_dataset'], update={'gpus': 1, - 'task.dataset': 'ase_db', 'optim.eval_every': 10, 'optim.max_epochs': 1, 'optim.batch_size': 4, 'logger':'tensorboard', # don't use wandb! # Train data 'dataset.train.src': 'train.db', + 'dataset.train.format': 'ase_db', 'dataset.train.a2g_args.r_energy': True, 'dataset.train.a2g_args.r_forces': True, # Test data - prediction only so no regression 'dataset.test.src': 'test.db', + 'dataset.test.format': 'ase_db', 'dataset.test.a2g_args.r_energy': False, 'dataset.test.a2g_args.r_forces': False, # val data 'dataset.val.src': 'val.db', + 'dataset.val.format': 'ase_db', 'dataset.val.a2g_args.r_energy': True, 'dataset.val.a2g_args.r_forces': True, }) diff --git a/docs/core/inference.md b/docs/core/inference.md index 919a301ac..fc717d71c 100644 --- a/docs/core/inference.md +++ b/docs/core/inference.md @@ -47,6 +47,12 @@ with ase.db.connect('full_data.db') as full_db: if 'tag' in atoms.info['key_value_pairs']: atoms.info['key_value_pairs']['tag'] = int(atoms.info['key_value_pairs']['tag']) + for key in atoms.info["key_value_pairs"]: + if atoms.info["key_value_pairs"][key] == "True": + atoms.info["key_value_pairs"][key] = True + elif atoms.info["key_value_pairs"][key] == "False": + atoms.info["key_value_pairs"][key] = False + subset_db.write(atoms, **atoms.info['key_value_pairs']) ``` @@ -78,11 +84,11 @@ yml = generate_yml_config(checkpoint_path, 'config.yml', 'dataset', 'slurm'], update={'amp': True, 'gpus': 1, - 'task.dataset': 'ase_db', 'task.prediction_dtype': 'float32', 'logger':'tensorboard', # don't use wandb! # Test data - prediction only so no regression 'dataset.test.src': 'data.db', + 'dataset.test.format': 'ase_db', 'dataset.test.a2g_args.r_energy': False, 'dataset.test.a2g_args.r_forces': False, 'dataset.test.select_args.selection': 'natoms>5,xc=PBE', diff --git a/docs/legacy_tutorials/OCP_Tutorial.md b/docs/legacy_tutorials/OCP_Tutorial.md index b4d1d4cfa..b05cef74b 100644 --- a/docs/legacy_tutorials/OCP_Tutorial.md +++ b/docs/legacy_tutorials/OCP_Tutorial.md @@ -302,7 +302,7 @@ dyn.run(fmax=0, steps=100) traj = ase.io.read("data/toy_c3h8_relax.traj", ":") # convert traj format to extxyz format (used by OC20 dataset) -columns = (['symbols','positions', 'move_mask', 'tags']) +columns = (['symbols','positions', 'move_mask', 'tags', 'forces']) with open('data/toy_c3h8_relax.extxyz','w') as f: extxyz.write_xyz(f, traj, columns=columns) ``` @@ -925,8 +925,8 @@ trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="s2ef", identifier="S2EF-example", run_dir=".", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1007,8 +1007,8 @@ pretrained_trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="s2ef", identifier="S2EF-val-example", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1211,8 +1211,8 @@ energy_trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="is2re", identifier="IS2RE-example", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1289,8 +1289,8 @@ pretrained_energy_trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="is2re", identifier="IS2RE-val-example", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1427,7 +1427,6 @@ task = { 'grad_input': 'atomic forces', 'train_on_free_atoms': True, 'eval_on_free_atoms': True, - 'relax_dataset': {"src": relax_dataset}, 'write_pos': True, 'relaxation_steps': 200, 'num_relaxation_batches': 1, @@ -1489,10 +1488,11 @@ optimizer = { 'force_coefficient': 100, } # Dataset -dataset = [ - {'src': train_src, 'normalize_labels': False}, # train set - {'src': val_src}, # val set (optional) -] +dataset = { + 'train': {'src': train_src, 'normalize_labels': False}, # train set + 'val': {'src': val_src}, # val set (optional) + 'relax': {"src": relax_dataset}, +} ``` @@ -1514,8 +1514,8 @@ trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="s2ef", identifier="is2rs-example", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1904,8 +1904,8 @@ trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="s2ef", identifier="S2EF-simple", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! @@ -1980,8 +1980,8 @@ trainer = OCPTrainer( dataset=dataset, optimizer=optimizer, outputs={}, - loss_fns={}, - eval_metrics={}, + loss_functions={}, + evaluation_metrics={}, name="s2ef", identifier="S2EF-gemnet-t", run_dir="./", # directory to save results if is_debug=False. Prediction files are saved here so be careful not to override! diff --git a/docs/legacy_tutorials/data_visualization.md b/docs/legacy_tutorials/data_visualization.md index 9d8aea4f2..ba26aefab 100644 --- a/docs/legacy_tutorials/data_visualization.md +++ b/docs/legacy_tutorials/data_visualization.md @@ -101,7 +101,7 @@ dyn.run(fmax=0, steps=100) traj = ase.io.read("data/toy_c3h8_relax.traj", ":") # convert traj format to extxyz format (used by OC20 dataset) -columns = (['symbols','positions', 'move_mask', 'tags']) +columns = (['symbols','positions', 'move_mask', 'tags', 'forces']) with open('data/toy_c3h8_relax.extxyz','w') as f: extxyz.write_xyz(f, traj, columns=columns) diff --git a/docs/tutorials/advanced/fine-tuning-in-python.md b/docs/tutorials/advanced/fine-tuning-in-python.md index afc810c20..568d92723 100644 --- a/docs/tutorials/advanced/fine-tuning-in-python.md +++ b/docs/tutorials/advanced/fine-tuning-in-python.md @@ -79,21 +79,23 @@ yml = generate_yml_config(checkpoint_path, 'config.yml', 'optim.loss_force', # the checkpoint setting causes an error 'dataset', 'test_dataset', 'val_dataset'], update={'gpus': 1, - 'task.dataset': 'ase_db', 'optim.eval_every': 10, 'optim.max_epochs': 1, 'optim.batch_size': 4, 'logger': 'tensorboard', # don't use wandb unless you already are logged in # Train data 'dataset.train.src': 'train.db', + 'dataset.train.format': 'ase_db', 'dataset.train.a2g_args.r_energy': True, 'dataset.train.a2g_args.r_forces': True, # Test data - prediction only so no regression 'dataset.test.src': 'test.db', + 'dataset.test.format': 'ase_db', 'dataset.test.a2g_args.r_energy': False, 'dataset.test.a2g_args.r_forces': False, # val data 'dataset.val.src': 'val.db', + 'dataset.val.format': 'ase_db', 'dataset.val.a2g_args.r_energy': True, 'dataset.val.a2g_args.r_forces': True, }) diff --git a/src/fairchem/core/common/relaxation/ase_utils.py b/src/fairchem/core/common/relaxation/ase_utils.py index 08efe08b6..152048d2c 100644 --- a/src/fairchem/core/common/relaxation/ase_utils.py +++ b/src/fairchem/core/common/relaxation/ase_utils.py @@ -38,7 +38,7 @@ def batch_to_atoms(batch): n_systems = batch.natoms.shape[0] natoms = batch.natoms.tolist() numbers = torch.split(batch.atomic_numbers, natoms) - fixed = torch.split(batch.fixed, natoms) + fixed = torch.split(batch.fixed.to(torch.bool), natoms) forces = torch.split(batch.force, natoms) positions = torch.split(batch.pos, natoms) tags = torch.split(batch.tags, natoms) diff --git a/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py b/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py index 43c833822..7ad3745ae 100644 --- a/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py +++ b/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py @@ -142,7 +142,7 @@ def run(self, fmax, steps): traj_fl = Path(self.traj_dir / f"{name}.traj_tmp", mode="w") traj_fl.rename(traj_fl.with_suffix(".traj")) - self.batch.y, self.batch.force = self.get_energy_and_forces( + self.batch.energy, self.batch.force = self.get_energy_and_forces( apply_constraint=False ) return self.batch @@ -211,7 +211,7 @@ def determine_step(dr): self.f0 = forces def write(self, energy, forces, update_mask) -> None: - self.batch.y, self.batch.force = energy, forces + self.batch.energy, self.batch.force = energy, forces atoms_objects = batch_to_atoms(self.batch) update_mask_ = torch.split(update_mask, self.batch.natoms.tolist()) for atm, traj, mask in zip(atoms_objects, self.trajectories, update_mask_): diff --git a/src/fairchem/core/common/tutorial_utils.py b/src/fairchem/core/common/tutorial_utils.py index 616537996..fadde8bcf 100644 --- a/src/fairchem/core/common/tutorial_utils.py +++ b/src/fairchem/core/common/tutorial_utils.py @@ -174,10 +174,6 @@ def nested_set(dic, keys, value): keys = _key.split(".") nested_set(config, keys, update[_key]) - # TODO : Do not rename keys in utils.py when reading a config - config["evaluation_metrics"] = config["eval_metrics"] - config.pop("eval_metrics") - out = dump(config) with open(yml, "wb") as f: f.write(out.encode("utf-8")) diff --git a/src/fairchem/core/trainers/base_trainer.py b/src/fairchem/core/trainers/base_trainer.py index 1daa5007f..12e2e61e7 100644 --- a/src/fairchem/core/trainers/base_trainer.py +++ b/src/fairchem/core/trainers/base_trainer.py @@ -159,10 +159,10 @@ def __init__( if len(dataset) > 2: self.config["test_dataset"] = dataset[2] elif isinstance(dataset, dict): - self.config["dataset"] = dataset.get("train", None) - self.config["val_dataset"] = dataset.get("val", None) - self.config["test_dataset"] = dataset.get("test", None) - self.config["relax_dataset"] = dataset.get("relax", None) + self.config["dataset"] = dataset.get("train", {}) + self.config["val_dataset"] = dataset.get("val", {}) + self.config["test_dataset"] = dataset.get("test", {}) + self.config["relax_dataset"] = dataset.get("relax", {}) else: self.config["dataset"] = dataset @@ -277,8 +277,19 @@ def load_datasets(self) -> None: self.val_loader = None self.test_loader = None + # Default all of the dataset portions to {} if + # they don't exist, or are null + if not self.config.get("dataset", None): + self.config["dataset"] = {} + if not self.config.get("val_dataset", None): + self.config["val_dataset"] = {} + if not self.config.get("test_dataset", None): + self.config["test_dataset"] = {} + if not self.config.get("relax_dataset", None): + self.config["relax_dataset"] = {} + # load train, val, test datasets - if self.config["dataset"].get("src", None): + if self.config["dataset"] and self.config["dataset"].get("src", None): logging.info( f"Loading dataset: {self.config['dataset'].get('format', 'lmdb')}" ) @@ -296,7 +307,7 @@ def load_datasets(self) -> None: self.train_sampler, ) - if self.config.get("val_dataset", None): + if self.config["val_dataset"]: if self.config["val_dataset"].get("use_train_settings", True): val_config = self.config["dataset"].copy() val_config.update(self.config["val_dataset"]) @@ -318,8 +329,13 @@ def load_datasets(self) -> None: self.val_sampler, ) - if self.config.get("test_dataset", None): - if self.config["test_dataset"].get("use_train_settings", True): + if self.config["test_dataset"]: + if ( + self.config["test_dataset"].get("use_train_settings", True) + and self.config[ + "dataset" + ] # if there's no training dataset, we have nothing to copy + ): test_config = self.config["dataset"].copy() test_config.update(self.config["test_dataset"]) else: @@ -340,30 +356,33 @@ def load_datasets(self) -> None: self.test_sampler, ) - if self.config.get("relax_dataset", None): - if self.config["relax_dataset"].get("use_train_settings", True): - relax_config = self.config["dataset"].copy() - relax_config.update(self.config["relax_dataset"]) - else: - relax_config = self.config["relax_dataset"] - - self.relax_dataset = registry.get_dataset_class( - relax_config.get("format", "lmdb") - )(relax_config) - self.relax_sampler = self.get_sampler( - self.relax_dataset, - self.config["optim"].get( - "eval_batch_size", self.config["optim"]["batch_size"] - ), - shuffle=False, - ) - self.relax_loader = self.get_dataloader( - self.relax_dataset, - self.relax_sampler, - ) + if self.config["relax_dataset"]: + if self.config["relax_dataset"].get("use_train_settings", True): + relax_config = self.config["dataset"].copy() + relax_config.update(self.config["relax_dataset"]) + else: + relax_config = self.config["relax_dataset"] + + self.relax_dataset = registry.get_dataset_class( + relax_config.get("format", "lmdb") + )(relax_config) + self.relax_sampler = self.get_sampler( + self.relax_dataset, + self.config["optim"].get( + "eval_batch_size", self.config["optim"]["batch_size"] + ), + shuffle=False, + ) + self.relax_loader = self.get_dataloader( + self.relax_dataset, + self.relax_sampler, + ) def load_task(self): # Normalizer for the dataset. + + # Is it troublesome that we assume any normalizer info is in train? What if there is no + # training dataset? What happens if we just specify a test normalizer = self.config["dataset"].get("transforms", {}).get("normalizer", {}) self.normalizers = {} if normalizer: @@ -388,16 +407,16 @@ def load_task(self): "outputs" ][target_name].get("level", "system") if "train_on_free_atoms" not in self.output_targets[subtarget]: - self.output_targets[subtarget]["train_on_free_atoms"] = ( - self.config[ - "outputs" - ][target_name].get("train_on_free_atoms", True) + self.output_targets[subtarget][ + "train_on_free_atoms" + ] = self.config["outputs"][target_name].get( + "train_on_free_atoms", True ) if "eval_on_free_atoms" not in self.output_targets[subtarget]: - self.output_targets[subtarget]["eval_on_free_atoms"] = ( - self.config[ - "outputs" - ][target_name].get("eval_on_free_atoms", True) + self.output_targets[subtarget][ + "eval_on_free_atoms" + ] = self.config["outputs"][target_name].get( + "eval_on_free_atoms", True ) # TODO: Assert that all targets, loss fn, metrics defined are consistent diff --git a/src/fairchem/core/trainers/ocp_trainer.py b/src/fairchem/core/trainers/ocp_trainer.py index 5b610aab6..9055d2d62 100644 --- a/src/fairchem/core/trainers/ocp_trainer.py +++ b/src/fairchem/core/trainers/ocp_trainer.py @@ -602,7 +602,7 @@ def run_relaxations(self, split="val"): s_idx += natoms target = { - "energy": relaxed_batch.y_relaxed, + "energy": relaxed_batch.energy, "positions": relaxed_batch.pos_relaxed[mask], "cell": relaxed_batch.cell, "pbc": torch.tensor([True, True, True]), @@ -610,7 +610,7 @@ def run_relaxations(self, split="val"): } prediction = { - "energy": relaxed_batch.y, + "energy": relaxed_batch.energy, "positions": relaxed_batch.pos[mask], "cell": relaxed_batch.cell, "pbc": torch.tensor([True, True, True]),