diff --git a/crystal_diffusion/models/score_network.py b/crystal_diffusion/models/score_network.py index 15c266fe..8bd5224f 100644 --- a/crystal_diffusion/models/score_network.py +++ b/crystal_diffusion/models/score_network.py @@ -121,7 +121,8 @@ def _forward_unchecked(self, batch: Dict[AnyStr, torch.Tensor]) -> torch.Tensor: 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 +138,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 diff --git a/examples/local_orion/config.yaml b/examples/local_orion/config.yaml deleted file mode 100644 index 6c9346c2..00000000 --- a/examples/local_orion/config.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# general -batch_size: 32 -optimizer: adam -loss: cross_entropy -max_epoch: 5 -exp_name: my_exp_1 -# set to null to avoid setting a seed (can speed up GPU computation, but -# results will not be reproducible) -seed: 1234 - -# architecture -num_classes: 10 -architecture: simple_mlp -hidden_dim: 'orion~uniform(32,256,discrete=True)' - -# early stopping -early_stopping: - metric: val_loss - mode: min - patience: 3 \ No newline at end of file diff --git a/examples/local_orion/diffusion/config_diffusion.yaml b/examples/local_orion/diffusion/config_diffusion.yaml new file mode 100644 index 00000000..48401594 --- /dev/null +++ b/examples/local_orion/diffusion/config_diffusion.yaml @@ -0,0 +1,45 @@ +# general +exp_name: example_experiment +run_name: run_debug_delete_me +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 diff --git a/examples/local_orion/orion_config.yaml b/examples/local_orion/diffusion/orion_config.yaml similarity index 100% rename from examples/local_orion/orion_config.yaml rename to examples/local_orion/diffusion/orion_config.yaml diff --git a/examples/local_orion/diffusion/run_orion.sh b/examples/local_orion/diffusion/run_orion.sh new file mode 100755 index 00000000..a84f2680 --- /dev/null +++ b/examples/local_orion/diffusion/run_orion.sh @@ -0,0 +1,16 @@ +export ORION_DB_ADDRESS='orion_db.pkl' +export ORION_DB_TYPE='pickleddb' + +CONFIG=config_diffusion.yaml +DATA_DIR=../../data/si_diffusion_small +PROCESSED_DATA=${DATA_DIR}/processed +DATA_WORK_DIR=./tmp_work_dir/ +OUTPUT=output + +orion -v hunt --config orion_config.yaml \ + python ../../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}/' diff --git a/examples/local_orion/run.sh b/examples/local_orion/run.sh deleted file mode 100644 index e8e14abc..00000000 --- a/examples/local_orion/run.sh +++ /dev/null @@ -1,8 +0,0 @@ -export MLFLOW_TRACKING_URI='mlruns' -export ORION_DB_ADDRESS='orion_db.pkl' -export ORION_DB_TYPE='pickleddb' - -orion -v hunt --config orion_config.yaml cd-train --data ../data \ - --config config.yaml --disable-progressbar \ - --output '{exp.working_dir}/{trial.id}/' \ - --log '{exp.working_dir}/{trial.id}/exp.log' diff --git a/tests/models/test_position_diffusion_lightning_model.py b/tests/models/test_position_diffusion_lightning_model.py index a193795e..b01beba2 100644 --- a/tests/models/test_position_diffusion_lightning_model.py +++ b/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, ) diff --git a/tests/models/test_score_network.py b/tests/models/test_score_network.py index 0fd1e7da..d992f096 100644 --- a/tests/models/test_score_network.py +++ b/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): diff --git a/tests/samplers/test_predictor_corrector_position_sampler.py b/tests/samplers/test_predictor_corrector_position_sampler.py index c92ce924..863a073e 100644 --- a/tests/samplers/test_predictor_corrector_position_sampler.py +++ b/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) diff --git a/tests/test_train_diffusion.py b/tests/test_train_diffusion.py index 92fe94b5..5e644bc9 100644 --- a/tests/test_train_diffusion.py +++ b/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)