Merge branch 'rlhf_2' of github.com:Delaunay/milabench into rlhf_2

mila-iqia · Aug 21, 2024 · 424ea99 · 424ea99
2 parents 63e78c4 + 70df378
commit 424ea99
Show file tree

Hide file tree

Showing 20 changed files with 121 additions and 91 deletions.
diff --git a/.gitignore b/.gitignore
@@ -32,7 +32,8 @@ scripts/article/cuda/
 scripts/article/xpu/
 
 dependencies/
-benchmarks/gflownet/gflownet
+benchmarks/recursiongfn/gflownet
+benchmarks/recursiongfn/logs/
 
 scripts/inventory.yaml
 output/

diff --git a/.pin/constraints-cuda-gnn.txt b/.pin/constraints-cuda-gnn.txt
diff --git a/.pin/constraints-cuda-torch.txt b/.pin/constraints-cuda-torch.txt
diff --git a/benchmarks/brax/requirements.cuda.txt b/benchmarks/brax/requirements.cuda.txt
diff --git a/benchmarks/diffusion/requirements.cuda.txt b/benchmarks/diffusion/requirements.cuda.txt
diff --git a/benchmarks/geo_gnn/bench/models.py b/benchmarks/geo_gnn/bench/models.py
@@ -14,8 +14,14 @@ def register_model(fn):
 def DimeNet(args, sample, **extras):
     # The directional message passing neural network (DimeNet) from the “Directional Message Passing for Molecular Graphs” paper. 
     # DimeNet transforms messages based on the angle between them in a rotation-equivariant fashion.
-    batch_size, out_channels = sample.y.shape
-
+
+    # PCQM4Mv2Subset: Data(x=[18, 9], edge_index=[2, 40], edge_attr=[40, 3], y=3.0476751256, pos=[18, 3], smiles='Cc1ccc([C@H]2[CH]c3cnccc3[N]C2=O)cc1')
+    # QM9:            Data(x=[5, 11], edge_index=[2, 8], edge_attr=[8, 4], y=[1, 19], pos=[5, 3], z=[5], smiles='[H]C([H])([H])[H]', name='gdb_1', idx=[1])
+    try:
+        batch_size, out_channels = sample.y.shape
+    except:
+        out_channels = 1
+
     return NS(
         category="3d",
         model=_DimeNet(

diff --git a/benchmarks/geo_gnn/main.py b/benchmarks/geo_gnn/main.py
@@ -107,10 +107,10 @@ def batch_size(x):
 
     observer = BenchObserver(batch_size_fn=batch_size)
 
-    train_dataset = PCQM4Mv2Subset(args.num_samples, args.root)
-    # train_dataset = QM9(args.root)
+    # train_dataset = PCQM4Mv2Subset(args.num_samples, args.root)
+    train_dataset = QM9(args.root)
 
-    sample = train_dataset[0]
+    sample = next(iter(train_dataset))
 
     info = models[args.model](args, 
                               sample=sample, 
@@ -137,8 +137,6 @@ def batch_size(x):
     device = accelerator.fetch_device(0)
     model = info.model.to(device)
 
-    print(model)
-
     criterion = nn.L1Loss()
 
     # set up optimizer
@@ -159,7 +157,13 @@ def batch_size(x):
             batch = batch.to(device)
 
             if args.use3d:
-                molecule_repr = model(z=batch.z, pos=batch.pos, batch=batch.batch)
+
+                if hasattr(batch, "z"):
+                    z = batch.z
+                else:
+                    z = batch.batch
+
+                molecule_repr = model(z=z, pos=batch.pos, batch=batch.batch)
             else:
                 molecule_repr = model(x=batch.x, batch=batch.batch, edge_index=batch.edge_index, batch_size=batch_size(batch))
 

diff --git a/benchmarks/geo_gnn/requirements-pre.cuda.txt b/benchmarks/geo_gnn/requirements-pre.cuda.txt
@@ -2,7 +2,7 @@
 # This file is autogenerated by pip-compile with Python 3.10
 # by the following command:
 #
-#    pip-compile --output-file=benchmarks/geo_gnn/requirements-pre.cuda.txt .pin/tmp-constraints-cuda-gnn.txt benchmarks/geo_gnn/requirements-pre.in
+#    pip-compile --output-file=benchmarks/geo_gnn/requirements-pre.cuda.txt .pin/tmp-constraints-cuda-geo_gnn.txt benchmarks/geo_gnn/requirements-pre.in
 #
 --extra-index-url https://download.pytorch.org/whl/cu121
 

diff --git a/benchmarks/geo_gnn/requirements.cuda.txt b/benchmarks/geo_gnn/requirements.cuda.txt
diff --git a/benchmarks/huggingface/requirements.cuda.txt b/benchmarks/huggingface/requirements.cuda.txt
diff --git a/benchmarks/huggingface/voirfile.py b/benchmarks/huggingface/voirfile.py
@@ -25,7 +25,7 @@ class Config:
     stop: int = 20
 
     # Number of seconds between each gpu poll
-    gpu_poll: int = 3
+    gpu_poll: int = 1
 
 
 @configurable