From 04580952375763757ea654169e064f821bdc7401 Mon Sep 17 00:00:00 2001 From: RaulPPealez Date: Fri, 26 Jan 2024 12:52:40 +0100 Subject: [PATCH 1/7] Update examples README --- examples/README.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/README.md b/examples/README.md index e8e2f20ba..59c993992 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,18 +1,20 @@ # Examples ## Training -We provide three example config files for the ET for training on QM9, MD17 and ANI1 respectively. To train on a QM9 target other than `energy_U0`, change the parameter `dataset_arg` in the QM9 config file. Changing the MD17 molecule to train on works analogously. To train an ET from scratch you can use the following code from the torchmd-net directory: + +You can reproduce any of the trainings in these folder using the torchmd-train utility: + ```bash -CUDA_VISIBLE_DEVICES=0,1 torchmd-train --conf examples/ET-{QM9,MD17,ANI1}.yaml +torchmd-train --conf [file].yaml ``` -Use the `CUDA_VISIBLE_DEVICES` environment variable to select which and how many GPUs you want to train on. The example above selects GPUs with indices 0 and 1. The training code will want to save checkpoints and config files in a directory called `logs/`, which you can change either in the config .yaml file or as an additional command line argument: `--log-dir path/to/log-dir`. +The training code will want to save checkpoints and config files in a directory called `logs/`, which you can change either in the config .yaml file or as an additional command line argument: `--log-dir path/to/log-dir`. ## Loading checkpoints You can access several pretrained checkpoint files under the following URLs: - equivariant Transformer pretrained on QM9 (U0): http://pub.htmd.org/et-qm9.zip - equivariant Transformer pretrained on MD17 (aspirin): http://pub.htmd.org/et-md17.zip - equivariant Transformer pretrained on ANI1: http://pub.htmd.org/et-ani1.zip -- invariant Transformer pretrained on ANI1: http://pub.htmd.org/t-ani1.zip + The checkpoints can be loaded using the `load_model` function in TorchMD-Net. Additional model arguments (e.g. turning on force prediction on top of energies) for inference can also be passed to the function. See the following example code for loading an ET pretrained on the ANI1 dataset: ```python From 90ff6695d554950be1bc4b489d1e2eeea2e7588b Mon Sep 17 00:00:00 2001 From: RaulPPealez Date: Fri, 26 Jan 2024 12:54:08 +0100 Subject: [PATCH 2/7] Add vector_cutoff to ET --- examples/ET-ANI1.yaml | 1 + examples/ET-MD17.yaml | 1 + examples/ET-QM9.yaml | 2 +- examples/ET-SPICE.yaml | 1 + 4 files changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/ET-ANI1.yaml b/examples/ET-ANI1.yaml index ff09a9a8c..a9b43f6f1 100644 --- a/examples/ET-ANI1.yaml +++ b/examples/ET-ANI1.yaml @@ -57,3 +57,4 @@ weight_decay: 0.0 box_vecs: null charge: false spin: false +vector_cutoff: true diff --git a/examples/ET-MD17.yaml b/examples/ET-MD17.yaml index cd7aaaa94..0e81d47e8 100644 --- a/examples/ET-MD17.yaml +++ b/examples/ET-MD17.yaml @@ -58,3 +58,4 @@ weight_decay: 0.0 box_vecs: null charge: false spin: false +vector_cutoff: true diff --git a/examples/ET-QM9.yaml b/examples/ET-QM9.yaml index 5a03d2f97..f2a9af787 100644 --- a/examples/ET-QM9.yaml +++ b/examples/ET-QM9.yaml @@ -59,4 +59,4 @@ box_vecs: null precision: 32 charge: false spin: false - +vector_cutoff: true diff --git a/examples/ET-SPICE.yaml b/examples/ET-SPICE.yaml index 84bfe6aad..25c7c2f20 100644 --- a/examples/ET-SPICE.yaml +++ b/examples/ET-SPICE.yaml @@ -58,3 +58,4 @@ weight_decay: 0.0 box_vecs: null charge: false spin: false +vector_cutoff: true From abc50b8cc791168ef949155ffef014f79a546fb5 Mon Sep 17 00:00:00 2001 From: RaulPPealez Date: Tue, 30 Jan 2024 12:24:44 +0100 Subject: [PATCH 3/7] Change seed --- tests/test_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_model.py b/tests/test_model.py index e1d22e051..044e6180b 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -196,7 +196,7 @@ def test_forward_output(model_name, output_model, overwrite_reference=False): @mark.parametrize("model_name", models.__all_models__) def test_gradients(model_name): - pl.seed_everything(1234) + pl.seed_everything(12345) precision = 64 output_model = "Scalar" # create model and sample batch From 41000c2bcf6d835b414f87bfdbdb9df1b61f8a45 Mon Sep 17 00:00:00 2001 From: RaulPPealez Date: Tue, 30 Jan 2024 12:41:07 +0100 Subject: [PATCH 4/7] Update reference results for ET test --- tests/expected.pkl | Bin 5944 -> 5944 bytes 1 file changed, 0 insertions(+), 0 deletions(-) diff --git a/tests/expected.pkl b/tests/expected.pkl index 0e22b0334f63600d4bb4c6f13d5be47625d06659..2efae2141f9b44582de8e6210b07817b99fb478f 100644 GIT binary patch delta 759 zcmY+C%`XE%7{)D4e6+r)o3;rd4$bcD%ycG=u80V6lQ@e%z`;et#l_9RVZv1qE-qvo zNGTE$Nwq%WEOFsR9HfYdw>z`lc26_EJkR?+?>l=GKZ@^97>$_N7PXPZscBJ>4@WAA zTs8)@2oavZl(?dpYxBsQF%z$@Qs0Y%&aStwGNP%gS43;hRXM~6p^Ta0uRSuQyqXlf z&DwkXw2}j9LQWDN-IH%#99vc5R?f4o#S-f1$JY+uyV~0DL_x)TdG1-!S&hDii z4+^JZCQvB#!mQW}hKYK;TvDDhAJ}!A**<9mLwY{bWYnVEolp5ZQtn&pyS-Rn7!yXY zU-ZXy{x6eW~OgFnPDiZnD5UCKUpyeV? zU(02UjK-$7SGZP0M>+a6O2btg#PH#xdW5pn2r(V(knz~p#K z#3+@qA@NPvF9UaqPvwEdKt+t`A5rS*WsC32v{i$gF)mc{&` XG{ADJQ*VcG+roqp|1Es-6HfgD=`zS# delta 760 zcmdm?w?l7(HzSX+rMZcjk%5V&v9aOgvtkOo$O1;2TNt$%x#2>FmXptjDJmlIO^pf} z6ZM!FAVA#PX#e%~hxSe86IY$AzboA*nI9FafE7Bsq{8kid7L zUtojHKmwMgrbh4}f{Sd{60t`ZV{8cYJhD|F0fWtpM4=%BPRu42X5grV7!69^V3A*9 Fb^xG#$D05E From 4a5d6e6674a609075ea5d4361c75c7cc11d3efa1 Mon Sep 17 00:00:00 2001 From: Guillem Simeon <55756547+guillemsimeon@users.noreply.github.com> Date: Tue, 30 Jan 2024 12:52:45 +0100 Subject: [PATCH 5/7] Update TensorNet-SPICE.yaml --- examples/TensorNet-SPICE.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/TensorNet-SPICE.yaml b/examples/TensorNet-SPICE.yaml index 4103b4418..0f3746c15 100644 --- a/examples/TensorNet-SPICE.yaml +++ b/examples/TensorNet-SPICE.yaml @@ -18,9 +18,9 @@ embed_files: null embedding_dimension: 128 energy_files: null equivariance_invariance_group: O(3) -y_weight: 0.5 +y_weight: 1.0 force_files: null -neg_dy_weight: 0.5 +neg_dy_weight: 10.0 gradient_clipping: 100.0 inference_batch_size: 16 load_model: null From 8efe48bf4663dec41be4603bc2e6e66fbc60aa8e Mon Sep 17 00:00:00 2001 From: RaulPPealez Date: Tue, 30 Jan 2024 13:16:38 +0100 Subject: [PATCH 6/7] Update SPICE version --- examples/ET-SPICE.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/ET-SPICE.yaml b/examples/ET-SPICE.yaml index 25c7c2f20..8205f18a8 100644 --- a/examples/ET-SPICE.yaml +++ b/examples/ET-SPICE.yaml @@ -8,7 +8,7 @@ cutoff_lower: 0.0 cutoff_upper: 10.0 dataset: SPICE dataset_arg: - version: 1.1.1 + version: 1.1.4 dataset_root: data derivative: true distance_influence: both From 7ae5286323d5a953126e6817916752da81043e71 Mon Sep 17 00:00:00 2001 From: Guillem Simeon <55756547+guillemsimeon@users.noreply.github.com> Date: Tue, 30 Jan 2024 13:25:35 +0100 Subject: [PATCH 7/7] Update SPICE version --- examples/TensorNet-SPICE.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/TensorNet-SPICE.yaml b/examples/TensorNet-SPICE.yaml index 0f3746c15..0c67409b0 100644 --- a/examples/TensorNet-SPICE.yaml +++ b/examples/TensorNet-SPICE.yaml @@ -7,7 +7,7 @@ cutoff_lower: 0.0 cutoff_upper: 10.0 dataset: SPICE dataset_arg: - version: 1.1.3 + version: 1.1.4 max_gradient: 50.94 dataset_root: ~/data derivative: true