Merge pull request #17 from RolnickLab/supernode

RolnickLab · Sep 16, 2022 · 4ca9b47 · 4ca9b47
2 parents f89f4ce + 9a75eb5
commit 4ca9b47
Show file tree

Hide file tree

Showing 23 changed files with 1,584 additions and 214 deletions.
diff --git a/configs/is2re/100k/dimenet_plus_plus/new_dpp.yml b/configs/is2re/100k/dimenet_plus_plus/new_dpp.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/100k/base.yml
+  - configs/is2re/100k/base.yml
 
 model:
   name: new_dimenetplusplus
@@ -16,9 +16,11 @@ model:
   use_pbc: True
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 1.

diff --git a/configs/is2re/100k/forcenet/new_forcenet.yml b/configs/is2re/100k/forcenet/new_forcenet.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/100k/base.yml
+  - configs/is2re/100k/base.yml
 
 model:
   name: new_forcenet
@@ -15,10 +15,13 @@ model:
   hidden_channels: 512
   decoder_hidden_channels: 512
   max_n: 3
+  use_pbc: True
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 8.

diff --git a/configs/is2re/100k/schnet/new_schnet.yml b/configs/is2re/100k/schnet/new_schnet.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/100k/base.yml
+  - configs/is2re/100k/base.yml
 
 model:
   name: new_schnet
@@ -12,9 +12,11 @@ model:
   regress_forces: False
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 1.

diff --git a/configs/is2re/10k/dimenet_plus_plus/new_dpp.yml b/configs/is2re/10k/dimenet_plus_plus/new_dpp.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/10k/base.yml
+  - configs/is2re/10k/base.yml
 
 model:
   name: new_dimenetplusplus
@@ -16,9 +16,11 @@ model:
   use_pbc: True
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 1.

diff --git a/configs/is2re/10k/forcenet/new_forcenet.yml b/configs/is2re/10k/forcenet/new_forcenet.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/10k/base.yml
+  - configs/is2re/10k/base.yml
 
 model:
   name: new_forcenet
@@ -15,10 +15,13 @@ model:
   hidden_channels: 512
   decoder_hidden_channels: 512
   max_n: 3
+  use_pbc: True
   # drlab attributes:
-  tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
-  phys_embeds: False # True
+  tag_hidden_channels: 64 # 64
+  pg_hidden_channels: 32 # 32 -> period & group embedding hidden channels
+  phys_embeds: True # True
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 8.

diff --git a/configs/is2re/10k/schnet/new_schnet.yml b/configs/is2re/10k/schnet/new_schnet.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/10k/base.yml
+  - configs/is2re/10k/base.yml
 
 model:
   name: new_schnet
@@ -12,9 +12,11 @@ model:
   regress_forces: False
   # drlab attributes:
   tag_hidden_channels: 0 # 32
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 1.

diff --git a/configs/is2re/all/dimenet_plus_plus/new_dpp.yml b/configs/is2re/all/dimenet_plus_plus/new_dpp.yml
@@ -16,9 +16,11 @@ model:
   use_pbc: True
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 4.

diff --git a/configs/is2re/all/forcenet/new_forcenet.yml b/configs/is2re/all/forcenet/new_forcenet.yml
@@ -15,11 +15,13 @@ model:
   hidden_channels: 512
   decoder_hidden_channels: 512
   max_n: 3
-  phys_embeds: True
+  use_pbc: True
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 8.

diff --git a/configs/is2re/all/schnet/new_schnet.yml b/configs/is2re/all/schnet/new_schnet.yml
@@ -1,5 +1,5 @@
 includes:
-- configs/is2re/all/base.yml
+  - configs/is2re/all/base.yml
 
 model:
   name: new_schnet
@@ -12,9 +12,11 @@ model:
   regress_forces: False
   # drlab attributes:
   tag_hidden_channels: 0 # 64
-  pg_hidden_channels: 0 # 32
+  pg_hidden_channels: 0 # 32 -> period & group embedding hidden channels
   phys_embeds: False # True
   phys_hidden_channels: 0
+  graph_rewiring: False # can be in {false, remove-tag-0, one-supernode-per-graph,
+                                    # one-supernode-per-atom-type, one-supernode-per-atom-type-dist}
 
 # *** Important note ***
 #   The total number of gpus used for this run was 4.

diff --git a/experiments.sh b/experiments.sh
@@ -1,23 +1,34 @@
-python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.phys_embeds=True --note='Baseline TAG + FIXED + PG embedding'" env=ocp
+python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --optim.lr_initial=0.0007 --model.graph_rewiring='one-supernode-per-graph' --note='smaller lr, pe + one-supernode-per-graph type'" env=ocp
 
-python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.phys_embeds=False --model.tag_hidden_channels=0 --note='Baseline PG embedding'" env=ocp
+python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/new_forcenet.yml --optim.lr_initial=0.0003 --model.graph_rewiring='one-supernode-per-graph' --note='smaller lr, pe + one-supernode-per-graph type'" env=ocp
 
-python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/dimenet_plus_plus/new_dpp.yml --model.phys_embeds=True --note='Baseline TAG + FIXED + PG embedding'" env=ocp
+python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/dimenet_plus_plus/new_dpp.yml --optim.lr_initial=0.00008 --model.graph_rewiring='one-supernode-per-graph' --note='smaller lr, pe + one-supernode-per-graph type'" env=ocp
 
-python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/dimenet_plus_plus/new_dpp.yml --model.phys_embeds=False --model.tag_hidden_channels=0 --model.phys_embeds=True --note='Baseline PG embedding'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --note='Baseline'" env=ocp
 
-#python sbatch.py gres="gpu:1" partition=long time=7:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/dimenet_plus_plus/new_dpp.yml --model.phys_embeds=True --model.tag_hidden_channels=0 --note='Baseline FIXED Embedding'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --optim.lr_initial=0.0007 --config-yml configs/is2re/all/schnet/new_schnet.yml --note='Baseline'" env=ocp
 
-#python sbatch.py gres="gpu:1" partition=long time=7:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/dimenet_plus_plus/new_dpp.yml --model.phys_embeds=True --note='Baseline TAG + FIXED Embedding'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/schnet.yml --note='True Baseline'" env=ocp
 
-#python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.phys_embeds=True --model.tag_hidden_channels=0 --note='Baseline FIXED embedding'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/forcenet.yml --note='True Baseline'" env=ocp
 
-#python sbatch.py gres="gpu:1" partition=long time=6:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.phys_embeds=True --note='Baseline TAG + FIXED embedding'" env=ocp
 
-# python sbatch.py gres="gpu:1" partition=long time=1:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.tag_hidden_channels=0 --model.hidden_channels=288 --note='Phys embed only, large hidden'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --model.graph_rewiring='remove-tag-0' --note='remove-tag-0'" env=ocp
 
-# python sbatch.py gres="gpu:1" partition=long time=3:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.tag_hidden_channels=0 --note='Phys embed only, same hidden'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --model.graph_rewiring='one-supernode-per-graph' --note='one-supernode-per-graph'" env=ocp
 
-# python sbatch.py gres="gpu:1" partition=long time=3:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.tag_hidden_channels=32 --model.hidden_channels=320 --note='Phys embed + tag, large hidden'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --model.graph_rewiring='one-supernode-per-atom-type' --note='one-supernode-per-atom-type'" env=ocp
 
-# python sbatch.py gres="gpu:1" partition=long time=3:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/10k/schnet/new_schnet.yml --model.tag_hidden_channels=32 --model.hidden_channels=256 --note='Phys embed + tag, same hidden'" env=ocp
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/schnet/new_schnet.yml --model.graph_rewiring='one-supernode-per-atom-type-dist' --note='one-supernode-per-atom-type-dist'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/new_forcenet.yml --model.graph_rewiring='remove-tag-0' --note='remove-tag-0'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/new_forcenet.yml --model.graph_rewiring='one-supernode-per-graph' --note='new one-supernode-per-graph'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/new_forcenet.yml --model.graph_rewiring='one-supernode-per-atom-type' --note='new  one supernode per atom type'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/forcenet/new_forcenet.yml --model.graph_rewiring='one-supernode-per-atom-type-dist' --note='new  one supernode per atom type dist'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/dimenet_plus_plus/new_dpp.yml --model.graph_rewiring='one-supernode-per-atom-type' --note='new  one supernode per atom type'" env=ocp
+
+# python sbatch.py gres="gpu:4" partition=long time=24:00:00 cpus=4 mem=32GB py_args="--mode train --config-yml configs/is2re/all/dimenet_plus_plus/new_dpp.yml --model.graph_rewiring='pe + one-supernode-per-graph' --note='new one supernode per graph'" env=ocp