Add config entry to test multiple location params.

experiments/mnist/mnist_digit_2_multiple_locs_scaled_14_linf_lognormal_gpu.yaml

@@ -0,0 +1,285 @@
+---
+__object__: src.explib.base.ExperimentCollection
+name: mnist_multiple_digits_poc
+experiments:
+  - &mnist_logNormal_linf_digit_0
+    __object__: src.explib.hyperopt.HyperoptExperiment
+    name: mnist_logNormal_linf_digit_2_loc_1
+    scheduler: &scheduler
+      __object__: ray.tune.schedulers.ASHAScheduler
+      max_t: 1000000
+      grace_period: 1000000
+      reduction_factor: 2
+    num_hyperopt_samples: &num_hyperopt_samples 3
+    gpus_per_trial: &gpus_per_trial 2
+    cpus_per_trial: &cpus_per_trial 0
+    tuner_params: &tuner_params
+      metric: val_loss
+      mode: min
+    device: &device cuda
+    trial_config:
+      logging:
+        images: true
+        "image_shape": [14, 14]
+      dataset: &dataset
+        __object__: src.explib.datasets.MnistSplit
+        scale: true
+        digit: 2
+        device: *device
+        scale_factor: 2
+      epochs: &epochs 200000
+      patience: &patience 150
+      batch_size: &batch_size
+        __eval__: tune.choice([16, 32])
+      optim_cfg: &optim
+        optimizer:
+          __class__: torch.optim.Adam
+        params:
+          lr:
+            __eval__: tune.loguniform(1e-4, 1e-3)
+          weight_decay: 0.0
+      model_cfg:
+        type:
+          __class__: &model src.veriflow.flows.NiceFlow
+        params:
+          soft_training: true
+          training_noise_prior:
+            __object__: pyro.distributions.Uniform
+            low:
+              __eval__: 1e-30 * torch.ones(1).to("cuda") #1e-20
+            high:
+                __eval__: 0.001 * torch.ones(1).to("cuda") #0.01
+          prior_scale: 5.0
+          coupling_layers: &coupling_layers
+            __eval__: tune.choice([i for i in range(3, 4)])
+          coupling_nn_layers: &coupling_nn_layers
+            __eval__: "tune.choice([[w] * l for l in [1] for w in [294]])"
+          nonlinearity: &nonlinearity
+            __eval__: tune.choice([torch.nn.ReLU()])
+          split_dim: 98
+          base_distribution:
+            __object__: src.veriflow.distributions.RadialDistribution
+            device: *device
+            p:
+              __eval__: math.inf
+            loc:
+              __eval__: torch.zeros(196).to("cuda")
+            norm_distribution:
+              __object__: pyro.distributions.LogNormal
+              loc:
+                __eval__: torch.ones(1).to("cuda")
+              scale:
+                __eval__: (0.5 * torch.ones(1)).to("cuda")
+          use_lu: false
+  - &mnist_logNormal_linf_digit_2_loc_1_2
+    name: mnist_logNormal_linf_digit_2_loc_1_2
+    __object__: src.explib.hyperopt.HyperoptExperiment
+    scheduler: *scheduler
+    num_hyperopt_samples: *num_hyperopt_samples
+    gpus_per_trial: *gpus_per_trial
+    cpus_per_trial: *cpus_per_trial
+    tuner_params: *tuner_params
+    device: *device
+    trial_config:
+      logging:
+        images: true
+        "image_shape": [ 14, 14 ]
+      dataset:
+        __object__: src.explib.datasets.MnistSplit
+        scale: true
+        digit: 2
+        device: *device
+        scale_factor: 2
+      epochs: *epochs
+      patience: *patience
+      batch_size: *batch_size
+      optim_cfg: *optim
+      model_cfg:
+        type:
+          __class__: *model
+        params:
+          soft_training: true
+          training_noise_prior:
+            __object__: pyro.distributions.Uniform
+            low:
+              __eval__: 1e-30 * torch.ones(1).to("cuda") #1e-20
+            high:
+              __eval__: 0.001 * torch.ones(1).to("cuda") #0.01
+          prior_scale: 5.0
+          coupling_layers: *coupling_layers
+          coupling_nn_layers: *coupling_nn_layers
+          nonlinearity: *nonlinearity
+          split_dim: 98
+          base_distribution:
+            __object__: src.veriflow.distributions.RadialDistribution
+            device: *device
+            p:
+              __eval__: math.inf
+            loc:
+              __eval__: torch.zeros(196).to("cuda")
+            norm_distribution:
+              __object__: pyro.distributions.LogNormal
+              loc:
+                __eval__: (1.2* torch.ones(1)).to("cuda")
+              scale:
+                __eval__: (0.5 * torch.ones(1)).to("cuda")
+          use_lu: false
+  - &mnist_logNormal_linf_digit_2_loc_1_4
+    name: mnist_logNormal_linf_digit_2_loc_1_4
+    __object__: src.explib.hyperopt.HyperoptExperiment
+    scheduler: *scheduler
+    num_hyperopt_samples: *num_hyperopt_samples
+    gpus_per_trial: *gpus_per_trial
+    cpus_per_trial: *cpus_per_trial
+    tuner_params: *tuner_params
+    device: *device
+    trial_config:
+      logging:
+        images: true
+        "image_shape": [ 14, 14 ]
+      dataset:
+        __object__: src.explib.datasets.MnistSplit
+        scale: true
+        digit: 2
+        device: *device
+        scale_factor: 2
+      epochs: *epochs
+      patience: *patience
+      batch_size: *batch_size
+      optim_cfg: *optim
+      model_cfg:
+        type:
+          __class__: *model
+        params:
+          soft_training: true
+          training_noise_prior:
+            __object__: pyro.distributions.Uniform
+            low:
+              __eval__: 1e-30 * torch.ones(1).to("cuda") #1e-20
+            high:
+              __eval__: 0.001 * torch.ones(1).to("cuda") #0.01
+          prior_scale: 5.0
+          coupling_layers: *coupling_layers
+          coupling_nn_layers: *coupling_nn_layers
+          nonlinearity: *nonlinearity
+          split_dim: 98
+          base_distribution:
+            __object__: src.veriflow.distributions.RadialDistribution
+            device: *device
+            p:
+              __eval__: math.inf
+            loc:
+              __eval__: torch.zeros(196).to("cuda")
+            norm_distribution:
+              __object__: pyro.distributions.LogNormal
+              loc:
+                __eval__: (1.4 * torch.ones(1)).to("cuda")
+              scale:
+                __eval__: (0.5 * torch.ones(1)).to("cuda")
+          use_lu: false
+  - &mnist_logNormal_linf_digit_2_loc_1_6
+    name: mnist_logNormal_linf_digit_2_loc_1_6
+    __object__: src.explib.hyperopt.HyperoptExperiment
+    scheduler: *scheduler
+    num_hyperopt_samples: *num_hyperopt_samples
+    gpus_per_trial: *gpus_per_trial
+    cpus_per_trial: *cpus_per_trial
+    tuner_params: *tuner_params
+    device: *device
+    trial_config:
+      logging:
+        images: true
+        "image_shape": [ 14, 14 ]
+      dataset:
+        __object__: src.explib.datasets.MnistSplit
+        scale: true
+        digit: 2
+        device: *device
+        scale_factor: 2
+      epochs: *epochs
+      patience: *patience
+      batch_size: *batch_size
+      optim_cfg: *optim
+      model_cfg:
+        type:
+          __class__: *model
+        params:
+          soft_training: true
+          training_noise_prior:
+            __object__: pyro.distributions.Uniform
+            low:
+              __eval__: 1e-30 * torch.ones(1).to("cuda") #1e-20
+            high:
+              __eval__: 0.001 * torch.ones(1).to("cuda") #0.01
+          prior_scale: 5.0
+          coupling_layers: *coupling_layers
+          coupling_nn_layers: *coupling_nn_layers
+          nonlinearity: *nonlinearity
+          split_dim: 98
+          base_distribution:
+            __object__: src.veriflow.distributions.RadialDistribution
+            device: *device
+            p:
+              __eval__: math.inf
+            loc:
+              __eval__: torch.zeros(196).to("cuda")
+            norm_distribution:
+              __object__: pyro.distributions.LogNormal
+              loc:
+                __eval__: (1.6 * torch.ones(1)).to("cuda")
+              scale:
+                __eval__: (0.5 * torch.ones(1)).to("cuda")
+          use_lu: false
+  - &mnist_logNormal_linf_digit_2_loc_1_8
+    name: mnist_logNormal_linf_digit_2_loc_1_8
+    __object__: src.explib.hyperopt.HyperoptExperiment
+    scheduler: *scheduler
+    num_hyperopt_samples: *num_hyperopt_samples
+    gpus_per_trial: *gpus_per_trial
+    cpus_per_trial: *cpus_per_trial
+    tuner_params: *tuner_params
+    device: *device
+    trial_config:
+      logging:
+        images: true
+        "image_shape": [ 14, 14 ]
+      dataset:
+        __object__: src.explib.datasets.MnistSplit
+        scale: true
+        digit: 2
+        device: *device
+        scale_factor: 2
+      epochs: *epochs
+      patience: *patience
+      batch_size: *batch_size
+      optim_cfg: *optim
+      model_cfg:
+        type:
+          __class__: *model
+        params:
+          soft_training: true
+          training_noise_prior:
+            __object__: pyro.distributions.Uniform
+            low:
+              __eval__: 1e-30 * torch.ones(1).to("cuda") #1e-20
+            high:
+              __eval__: 0.001 * torch.ones(1).to("cuda") #0.01
+          prior_scale: 5.0
+          coupling_layers: *coupling_layers
+          coupling_nn_layers: *coupling_nn_layers
+          nonlinearity: *nonlinearity
+          split_dim: 98
+          base_distribution:
+            __object__: src.veriflow.distributions.RadialDistribution
+            device: *device
+            p:
+              __eval__: math.inf
+            loc:
+              __eval__: torch.zeros(196).to("cuda")
+            norm_distribution:
+              __object__: pyro.distributions.LogNormal
+              loc:
+                __eval__: (1.8 * torch.ones(1)).to("cuda")
+              scale:
+                __eval__: (0.5 * torch.ones(1)).to("cuda")
+          use_lu: false