Add config files and links to pretrained models

Author: dav0dea

examples/pretrained/ANI1-equivariant_transformer/hparams.yaml renamed to examples/ET-ANI1.yaml

@@ -1,60 +1,57 @@
-load_model: null
-conf: null
-num_epochs: 3000
-batch_size: 512
-inference_batch_size: 512
-lr: 0.0007
-lr_patience: 5
-lr_min: 1.0e-07
-lr_factor: 0.5
-lr_warmup_steps: 10000
-early_stopping_patience: 500
-weight_decay: 0.0
-ema_alpha_y: 1.0
-ema_alpha_dy: 1.0
-ngpus: -1
-num_nodes: 1
-precision: 32
-log_dir: /shared/philipp/geometric/logs/equivariant_transformer-ANI1/
-splits: null
-train_size: 0.8
-val_size: 0.05
-test_size: null
-test_interval: 10
-save_interval: 10
-seed: 10
-distributed_backend: ddp
-num_workers: 6
-redirect: true
+activation: silu
+aggr: add
+atom_filter: -1
+attn_activation: silu
+batch_size: 2048
+coord_files: null
+cutoff_lower: 0.0
+cutoff_upper: 5.0
 dataset: ANI1
-dataset_root: ~/data
 dataset_arg: null
-coord_files: null
+dataset_root: ~/data
+derivative: false
+distance_influence: both
+distributed_backend: ddp
+early_stopping_patience: 500
+ema_alpha_dy: 1.0
+ema_alpha_y: 1.0
 embed_files: null
+embedding_dimension: 128
 energy_files: null
-force_files: null
 energy_weight: 1.0
+force_files: null
 force_weight: 1.0
+inference_batch_size: 2048
+load_model: null
+log_dir: logs/
+lr: 0.0007
+lr_factor: 0.5
+lr_min: 1.0e-07
+lr_patience: 5
+lr_warmup_steps: 10000
+max_num_neighbors: 32
+max_z: 100
 model: equivariant-transformer
+neighbor_embedding: true
+ngpus: -1
+num_epochs: 3000
+num_heads: 8
+num_layers: 6
+num_nodes: 1
+num_rbf: 32
+num_workers: 6
 output_model: Scalar
+precision: 32
 prior_model: Atomref
-embedding_dimension: 128
-num_layers: 6
-num_rbf: 64
-activation: silu
 rbf_type: expnorm
-trainable_rbf: false
-neighbor_embedding: true
-distance_influence: both
-attn_activation: silu
-num_heads: 8
-derivative: false
-cutoff_lower: 0.0
-cutoff_upper: 5.0
-atom_filter: -1
-max_z: 100
-max_num_neighbors: 32
-standardize: false
+redirect: false
 reduce_op: add
-prior_args:
-  max_z: 100
+save_interval: 10
+splits: null
+standardize: false
+test_interval: 10
+test_size: null
+train_size: 0.8
+trainable_rbf: false
+val_size: 0.05
+weight_decay: 0.0

examples/example.yaml renamed to examples/ET-MD17.yaml

@@ -2,56 +2,56 @@ activation: silu
 aggr: add
 atom_filter: -1
 attn_activation: silu
-batch_size: 16
+batch_size: 8
 coord_files: null
 cutoff_lower: 0.0
 cutoff_upper: 5.0
-dataset: QM9
-dataset_arg: energy_U0
-dataset_root: /home/philipp/Documents/data/qm9
-derivative: false
+dataset: MD17
+dataset_arg: aspirin
+dataset_root: ~/data
+derivative: true
 distance_influence: both
-distributed_backend: ddp_cpu
-early_stopping_patience: 30
+distributed_backend: ddp
+early_stopping_patience: 300
 ema_alpha_dy: 1.0
-ema_alpha_y: 1.0
+ema_alpha_y: 0.05
 embed_files: null
 embedding_dimension: 128
 energy_files: null
-energy_weight: 1.0
+energy_weight: 0.2
 force_files: null
-force_weight: 1.0
-inference_batch_size: 16
+force_weight: 0.8
+inference_batch_size: 64
 load_model: null
-log_dir: /tmp/logs
-lr: 0.0001
+log_dir: logs/
+lr: 0.001
 lr_factor: 0.8
-lr_min: 1.0e-06
-lr_patience: 10
-lr_warmup_steps: 0
+lr_min: 1.0e-07
+lr_patience: 30
+lr_warmup_steps: 1000
 max_num_neighbors: 32
 max_z: 100
+model: equivariant-transformer
 neighbor_embedding: true
-ngpus: 0
-num_epochs: 300
+ngpus: -1
+num_epochs: 3000
 num_heads: 8
 num_layers: 6
 num_nodes: 1
-num_rbf: 20
-num_workers: 4
+num_rbf: 32
+num_workers: 6
 output_model: Scalar
 precision: 32
-prior_model: Atomref
+prior_model: null
 rbf_type: expnorm
 redirect: false
 reduce_op: add
-save_interval: 3
-seed: 1
+save_interval: 10
 splits: null
-standardize: false
-test_interval: 3
+standardize: true
+test_interval: 10
 test_size: null
-train_size: 110000
+train_size: 950
 trainable_rbf: false
-val_size: 10000
+val_size: 50
 weight_decay: 0.0

examples/pretrained/ANI1-transformer/hparams.yaml renamed to examples/ET-QM9.yaml

@@ -1,60 +1,57 @@
+activation: silu
+aggr: add
+atom_filter: -1
+attn_activation: silu
+batch_size: 128
+coord_files: null
+cutoff_lower: 0.0
+cutoff_upper: 5.0
+dataset: QM9
+dataset_arg: energy_U0
+dataset_root: ~/data
+derivative: false
+distance_influence: both
+distributed_backend: ddp
+early_stopping_patience: 150
+ema_alpha_dy: 1.0
+ema_alpha_y: 1.0
+embed_files: null
+embedding_dimension: 256
+energy_files: null
+energy_weight: 1.0
+force_files: null
+force_weight: 1.0
+inference_batch_size: 128
 load_model: null
-conf: null
-num_epochs: 3000
-batch_size: 300
-inference_batch_size: 300
-lr: 0.01
-lr_patience: 5
+log_dir: logs/
+lr: 0.0004
+lr_factor: 0.8
 lr_min: 1.0e-07
-lr_factor: 0.5
+lr_patience: 15
 lr_warmup_steps: 10000
-early_stopping_patience: 500
-weight_decay: 0.0
-ema_alpha_y: 1.0
-ema_alpha_dy: 1.0
+max_num_neighbors: 32
+max_z: 100
+model: equivariant-transformer
+neighbor_embedding: true
 ngpus: -1
+num_epochs: 3000
+num_heads: 8
+num_layers: 8
 num_nodes: 1
-precision: 32
-log_dir: /shared/philipp/geometric/logs/transformer-ANI1/
-splits: null
-train_size: 0.8
-val_size: 0.05
-test_size: null
-test_interval: 10
-save_interval: 10
-seed: 10
-distributed_backend: ddp
+num_rbf: 64
 num_workers: 6
-redirect: true
-dataset: ANI1
-dataset_root: ~/data
-dataset_arg: null
-coord_files: null
-embed_files: null
-energy_files: null
-force_files: null
-energy_weight: 1.0
-force_weight: 1.0
-model: transformer
 output_model: Scalar
+precision: 32
 prior_model: Atomref
-embedding_dimension: 256
-num_layers: 8
-num_rbf: 64
-activation: silu
 rbf_type: expnorm
-trainable_rbf: false
-neighbor_embedding: true
-distance_influence: both
-attn_activation: silu
-num_heads: 8
-derivative: false
-cutoff_lower: 0.0
-cutoff_upper: 5.0
-atom_filter: -1
-max_z: 100
-max_num_neighbors: 32
-standardize: false
+redirect: false
 reduce_op: add
-prior_args:
-  max_z: 100
+save_interval: 10
+splits: null
+standardize: false
+test_interval: 10
+test_size: null
+train_size: 110000
+trainable_rbf: false
+val_size: 10000
+weight_decay: 0.0

examples/README.md

@@ -0,0 +1,36 @@
+# 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:
+```bash
+CUDA_VISIBLE_DEVICES=0,1 python scripts/train.py --conf examples/ET-{QM9,MD17,ANI1}.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`.
+
+## 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 torchmdnet.models.model import load_model
+model = load_model("ANI1-equivariant_transformer/epoch=209-val_loss=0.0003-test_loss=0.0093.ckpt", derivative=True)
+```
+The following example shows how to run inference on the model checkpoint. For single molecules, you just have to pass atomic numbers and position tensors, to evaluate the model on multiple molecules simultaneously, also include a batch tensor, which contains the molecule index of each atom.
+```python
+import torch
+
+# single molecule
+z = torch.tensor([1, 1, 8], dtype=torch.long)
+pos = torch.rand(3, 3)
+energy, forces = model(z, pos)
+
+# multiple molecules
+z = torch.tensor([1, 1, 8, 1, 1, 8], dtype=torch.long)
+pos = torch.rand(6, 3)
+batch = torch.tensor([0, 0, 0, 1, 1, 1], dtype=torch.long)
+energies, forces = model(z, pos, batch)
+```