Merge pull request #22 from mila-iqia/orion_for_diffusion_training

orion script and rewrite how layers are defined in MLP score network

crystal_diffusion/models/score_network.py

@@ -5,7 +5,7 @@
 periodic unit cell.
 """
 from dataclasses import dataclass
-from typing import AnyStr, Dict, List
+from typing import AnyStr, Dict
 
 import torch
 from torch import nn
@@ -120,8 +120,10 @@ def _forward_unchecked(self, batch: Dict[AnyStr, torch.Tensor]) -> torch.Tensor:
 @dataclass(kw_only=True)
 class MLPScoreNetworkParameters(BaseScoreNetworkParameters):
     """Specific Hyper-parameters for MLP score networks."""
+
     number_of_atoms: int  # the number of atoms in a configuration.
-    hidden_dimensions: List[int]  # dimensions of the hidden layers. Length of array determines number of layers.
+    n_hidden_dimensions: int  # the number of hidden layers.
+    hidden_dimensions_size: int  # the dimensions of the hidden layers.
 
 
 class MLPScoreNetwork(ScoreNetwork):
@@ -137,7 +139,7 @@ def __init__(self, hyper_params: MLPScoreNetworkParameters):
             hyper_params (dict): hyper parameters from the config file.
         """
         super(MLPScoreNetwork, self).__init__(hyper_params)
-        hidden_dimensions = hyper_params.hidden_dimensions
+        hidden_dimensions = [hyper_params.hidden_dimensions_size] * hyper_params.n_hidden_dimensions
         self._natoms = hyper_params.number_of_atoms
 
         output_dimension = self.spatial_dimension * self._natoms

examples/local_orion/diffusion/config_diffusion.yaml

@@ -0,0 +1,44 @@
+# general
+exp_name: example_experiment
+max_epoch: 2
+log_every_n_steps: 1
+# fast_dev_run: True
+# set to null to avoid setting a seed (can speed up GPU computation, but
+# results will not be reproducible)
+seed: 1234
+
+# data
+data:
+  batch_size: 128
+  num_workers: 4
+  max_atom: 64
+
+# architecture
+spatial_dimension: 3
+model:
+  score_network:
+    n_hidden_dimensions: 'orion~choices([1, 2, 3, 4])'
+    hidden_dimensions_size: 'orion~choices([256, 512, 1024, 2048])'
+  noise:
+    total_time_steps: 'orion~uniform(2, 20, discrete=True)'
+    sigma_min: 'orion~choices([0.001, 0.005, 0.01])'
+    sigma_max: 'orion~choices([0.1, 0.5, 0.75])'
+
+# optimizer
+optimizer:
+  name: adam
+  learning_rate: 'orion~loguniform(1e-6, 1e-3)'
+
+# early stopping
+early_stopping:
+  metric: val_loss
+  mode: min
+  patience: 100
+
+model_checkpoint:
+  monitor: val_loss
+  mode: min
+
+logging:
+  - tensorboard
+  - comet

examples/local_orion/diffusion/run_orion.sh

@@ -0,0 +1,16 @@
+export ORION_DB_ADDRESS='orion_db.pkl'
+export ORION_DB_TYPE='pickleddb'
+
+ROOT_DIR=../../../
+CONFIG=config_diffusion.yaml
+DATA_DIR=${ROOT_DIR}/data/si_diffusion_small
+PROCESSED_DATA=${DATA_DIR}/processed
+DATA_WORK_DIR=./tmp_work_dir/
+
+orion -v hunt --config orion_config.yaml \
+    python ${ROOT_DIR}/crystal_diffusion/train_diffusion.py \
+    --config $CONFIG \
+    --data $DATA_DIR \
+    --processed_datadir $PROCESSED_DATA \
+    --dataset_working_dir $DATA_WORK_DIR \
+    --output '{exp.working_dir}/{trial.id}/'

tests/models/test_position_diffusion_lightning_model.py

@@ -69,7 +69,8 @@ def number_of_atoms(self):
     def hyper_params(self, number_of_atoms, spatial_dimension):
         score_network_parameters = MLPScoreNetworkParameters(
             number_of_atoms=number_of_atoms,
-            hidden_dimensions=[4, 8, 16],
+            n_hidden_dimensions=3,
+            hidden_dimensions_size=8,
             spatial_dimension=spatial_dimension,
         )
 

tests/models/test_score_network.py

@@ -74,7 +74,8 @@ def test_check_batch_bad(self, base_score_network, bad_batch):
 
 
 @pytest.mark.parametrize("spatial_dimension", [2, 3])
-@pytest.mark.parametrize("hidden_dimensions", [[16], [8, 16], [8, 16, 32]])
+@pytest.mark.parametrize("n_hidden_dimensions", [1, 2, 3])
+@pytest.mark.parametrize("hidden_dimensions_size", [8, 16])
 class TestMLPScoreNetwork:
 
     @pytest.fixture()
@@ -102,10 +103,11 @@ def bad_batch(self, batch_size, number_of_atoms, spatial_dimension):
         return {ScoreNetwork.position_key: positions, ScoreNetwork.timestep_key: times}
 
     @pytest.fixture()
-    def score_network(self, number_of_atoms, spatial_dimension, hidden_dimensions):
+    def score_network(self, number_of_atoms, spatial_dimension, n_hidden_dimensions, hidden_dimensions_size):
         hyper_params = MLPScoreNetworkParameters(spatial_dimension=spatial_dimension,
                                                  number_of_atoms=number_of_atoms,
-                                                 hidden_dimensions=hidden_dimensions)
+                                                 n_hidden_dimensions=n_hidden_dimensions,
+                                                 hidden_dimensions_size=hidden_dimensions_size)
         return MLPScoreNetwork(hyper_params)
 
     def test_check_batch_bad(self, score_network, bad_batch):

tests/samplers/test_predictor_corrector_position_sampler.py

@@ -85,7 +85,8 @@ def test_sample(self, sampler, number_of_samples, expected_samples):
 @pytest.mark.parametrize("total_time_steps", [1, 5, 10])
 @pytest.mark.parametrize("number_of_corrector_steps", [0, 1, 2])
 @pytest.mark.parametrize("spatial_dimension", [2, 3])
-@pytest.mark.parametrize("hidden_dimensions", [[8, 16, 32]])
+@pytest.mark.parametrize("n_hidden_dimensions", [3])
+@pytest.mark.parametrize("hidden_dimensions_size", [16])
 @pytest.mark.parametrize("number_of_atoms", [8])
 @pytest.mark.parametrize("time_delta", [0.1])
 @pytest.mark.parametrize("sigma_min", [0.15])
@@ -94,12 +95,13 @@ def test_sample(self, sampler, number_of_samples, expected_samples):
 class TestAnnealedLangevinDynamics:
     @pytest.fixture()
     def sigma_normalized_score_network(
-        self, number_of_atoms, spatial_dimension, hidden_dimensions
+        self, number_of_atoms, spatial_dimension, n_hidden_dimensions, hidden_dimensions_size
     ):
         hyper_params = MLPScoreNetworkParameters(
             spatial_dimension=spatial_dimension,
             number_of_atoms=number_of_atoms,
-            hidden_dimensions=hidden_dimensions,
+            n_hidden_dimensions=n_hidden_dimensions,
+            hidden_dimensions_size=hidden_dimensions_size
         )
         return MLPScoreNetwork(hyper_params)
 

tests/test_train_diffusion.py

@@ -62,7 +62,8 @@ def get_fake_configuration_dataframe(number_of_configurations, number_of_atoms):
 def get_config(number_of_atoms: int, max_epoch: int):
     data_config = dict(batch_size=4, num_workers=0, max_atom=number_of_atoms)
 
-    model_config = dict(score_network={'hidden_dimensions': [16, 16]},
+    model_config = dict(score_network={'n_hidden_dimensions': 2,
+                                       'hidden_dimensions_size': 16},
                         noise={'total_time_steps': 10})
 
     optimizer_config = dict(name='adam', learning_rate=0.001)