Dev rh pythae

clinicadl/random_search/random_search.py

@@ -2,21 +2,48 @@
 Launch a random network training.
 """
 from pathlib import Path
-
-from clinicadl.random_search.random_search_utils import get_space_dict, random_sampling
-from clinicadl.train import train
+from random import sample
 
 
 def launch_search(launch_directory: Path, job_name):
+    from clinicadl.random_search.random_search_classification_utils import (
+        classification_random_sampling,
+        get_classification_space_dict,
+    )
+    from clinicadl.train import train
+
     if not (launch_directory / "random_search.toml").is_file():
         raise FileNotFoundError(
             f"TOML file 'random_search.toml' must be written in directory: {launch_directory}."
         )
-    space_options = get_space_dict(launch_directory)
-    options = random_sampling(space_options)
+    space_options = get_classification_space_dict(launch_directory)
+    options = classification_random_sampling(space_options)
 
     maps_directory = launch_directory / job_name
     split = options.pop("split")
     options["architecture"] = "RandomArchitecture"
 
     train(maps_directory, options, split)
+
+
+def launch_vae_search(launch_directory, job_name):
+    from clinicadl.random_search.random_search_vae_utils import (
+        get_vae_space_dict,
+        vae_random_sampling,
+    )
+    from clinicadl.utils.maps_manager import MapsManager
+
+    space_options = get_vae_space_dict(launch_directory)
+    parameters = vae_random_sampling(space_options)
+    parameters["architecture"] = "pythae_VAE"
+    print("Parameters:", parameters)
+
+    # Select 3 splits randomly
+    split_list = sample(range(6), 3)
+    print("Split list:", split_list)
+
+    # initialise maps
+    maps_dir = launch_directory / job_name
+    maps_manager = MapsManager(maps_dir, parameters, verbose="info")
+    # launch training procedure for Pythae
+    maps_manager.train_pythae(split_list=split_list)

clinicadl/random_search/random_search_utils.py → clinicadl/random_search/random_search_classification_utils.py

@@ -10,7 +10,7 @@
 from clinicadl.utils.preprocessing import read_preprocessing
 
 
-def get_space_dict(launch_directory: Path) -> Dict[str, Any]:
+def get_classification_space_dict(launch_directory: str) -> Dict[str, Any]:
     """Transforms the TOML dictionary in one dimension dictionary."""
     toml_path = launch_directory / "random_search.toml"
     toml_options = toml.load(toml_path)
@@ -98,7 +98,7 @@ def sampling_fn(value, sampling_type: str):
             return value
 
 
-def random_sampling(rs_options: Dict[str, Any]) -> Dict[str, Any]:
+def classification_random_sampling(rs_options: Dict[str, Any]) -> Dict[str, Any]:
     """
     Samples all the hyperparameters of the model.
     Args:

clinicadl/random_search/random_search_cli.py

@@ -4,12 +4,14 @@
 
 
 @click.command("random-search", no_args_is_help=True)
+@click.argument("task", type=click.Choice(["classification", "vae-architecture"]))
 @click.argument(
     "launch_directory",
     type=click.Path(exists=True, path_type=Path),
 )
 @click.argument("name", type=str)
 def cli(
+    task,
     launch_directory,
     name,
 ):
@@ -19,9 +21,15 @@ def cli(
 
     NAME is the name of the output folder containing the experiment.
     """
-    from .random_search import launch_search
 
-    launch_search(launch_directory, name)
+    if task == "classification":
+        from .random_search import launch_search
+
+        launch_search(task, launch_directory, name)
+    elif task == "vae-architecture":
+        from .random_search import launch_vae_search
+
+        launch_vae_search(launch_directory, name)
 
 
 if __name__ == "__main__":

clinicadl/random_search/random_search_vae_utils.py

@@ -0,0 +1,137 @@
+import random
+from os import path
+from typing import Any, Dict, Tuple
+
+import toml
+
+from clinicadl.random_search.random_search_classification_utils import sampling_fn
+from clinicadl.train.train_utils import build_train_dict
+from clinicadl.utils.exceptions import ClinicaDLConfigurationError
+from clinicadl.utils.preprocessing import read_preprocessing
+
+
+def get_vae_space_dict(launch_directory):
+    """
+    Takes a launch directory with a "random_search.toml" file with all the parameters to explore.
+    Return a parameters dictionnary randomly sampled
+    """
+
+    toml_path = path.join(launch_directory, "random_search.toml")
+    if not path.exists(toml_path):
+        raise FileNotFoundError(
+            f"TOML file 'random_search.toml' must be written in directory {launch_directory}."
+        )
+
+    # load TOML file and create space dict
+    toml_options = toml.load(toml_path)
+    space_dict = dict()
+
+    # check and read TOML
+    if "Random_Search" not in toml_options:
+        raise ClinicaDLConfigurationError(
+            "Category 'Random_Search' must be defined in the random_search.toml file. "
+            "All random search arguments AND options must be defined in this category."
+        )
+
+    for key in toml_options["Random_Search"]:
+        space_dict[key] = toml_options["Random_Search"][key]
+
+    # Check presence of mandatory arguments
+    mandatory_arguments = [
+        "network_task",
+        "tsv_path",
+        "caps_directory",
+        "preprocessing_json",
+        "first_layer_channels",
+        "n_conv_encoder",
+        "feature_size",
+        "latent_space_size",
+        "n_conv_decoder",
+        "last_layer_channels",
+        "last_layer_conv",
+    ]
+
+    for argument in mandatory_arguments:
+        if argument not in space_dict:
+            raise ClinicaDLConfigurationError(
+                f"The argument {argument} must be specified in the random_search.toml file (Random_Search category)."
+            )
+
+    # Make training parameter dict
+    train_default = build_train_dict(toml_path, space_dict["network_task"])
+
+    # Mode and preprocessing
+    preprocessing_json = path.join(
+        space_dict["caps_directory"],
+        "tensor_extraction",
+        space_dict.pop("preprocessing_json"),
+    )
+
+    preprocessing_dict = read_preprocessing(preprocessing_json)
+    train_default["preprocessing_dict"] = preprocessing_dict
+    train_default["mode"] = preprocessing_dict["mode"]
+
+    # Add the other default parameters to the dictionnary
+    # space_dict.update(train_default)
+    for k, v in train_default.items():
+        if k not in space_dict:
+            space_dict[k] = v
+    return space_dict
+
+
+def vae_random_sampling(space_dict):
+    # Create parameters dict
+    parameters = dict()
+
+    sampling_vae_dict = {
+        "accumulation_steps": "fixed",
+        "baseline": "fixed",
+        "batch_size": "fixed",
+        "caps_directory": "fixed",
+        "channels_limit": "fixed",
+        "compensation": "fixed",
+        "data_augmentation": "fixed",
+        "deterministic": "fixed",
+        "diagnoses": "fixed",
+        "epochs": "fixed",
+        "evaluation_steps": "fixed",
+        "gpu": "fixed",
+        "label": "fixed",
+        "learning_rate": "fixed",
+        "mode": "fixed",
+        "multi_cohort": "fixed",
+        "n_splits": "fixed",
+        "n_proc": "fixed",
+        "network_task": "fixed",
+        "normalize": "fixed",
+        "optimizer": "fixed",
+        "patience": "fixed",
+        "preprocessing_dict": "fixed",
+        "sampler": "fixed",
+        "seed": "fixed",
+        "selection_metrics": "fixed",
+        "size_reduction": "fixed",
+        "size_reduction_factor": "fixed",
+        "split": "fixed",
+        "tolerance": "fixed",
+        "transfer_path": "fixed",
+        "transfer_selection_metric": "fixed",
+        "tsv_path": "fixed",
+        "wd_bool": "fixed",
+        "weight_decay": "fixed",
+        # VAE architecture
+        "first_layer_channels": "choice",
+        "n_conv_encoder": "randint",
+        "feature_size": "choice",
+        "latent_space_size": "choice",
+        "n_conv_decoder": "randint",
+        "last_layer_channels": "choice",
+        "last_layer_conv": "choice",
+    }
+
+    for name, sampling_type in sampling_vae_dict.items():
+        if name in space_dict:
+            sampled_value = sampling_fn(space_dict[name], sampling_type)
+            parameters[name] = sampled_value
+
+    return parameters

clinicadl/resources/config/train_config.toml

@@ -10,10 +10,13 @@ multi_network = false
 dropout = 0.0 # between 0 and 1
 # architecture.n_layers = 4
 # VAE
+first_layer_channels = 32
+n_conv_encoder = 4
+feature_size = 0
 latent_space_size = 128
-feature_size = 1024
-n_conv = 4
-io_layer_channels = 8
+n_conv_decoder = 4
+last_layer_channels = 32
+last_layer_conv = false
 recons_weight = 1
 kl_weight = 1
 normalization = "batch"
@@ -34,6 +37,86 @@ loss = "MSELoss"
 selection_metrics = ["loss"]
 loss = "MSELoss"
 
+[Pythae]
+# Beta VAE
+beta=5
+# Linear Normalizing Flow VAE
+flows=['Planar', 'Radial', 'Planar']
+# Inverse Autoregressive Flows
+n_made_blocks=2
+n_hidden_in_made=3
+hidden_size=128
+# Disentangled Beta VAE
+#beta=5
+C=30.0
+warmup_epoch=25
+# Factor VAE
+gamma=10
+# Beta TC VAE
+#beta=2.
+alpha=1
+#gamma=1
+# MS SSIM VAE
+#beta=1e-2
+window_size=3
+# Info VAE
+kernel_choice='imq'
+#alpha=-2
+lbd=10
+kernel_bandwidth=1
+# Wasserstein Autoencoder
+#kernel_choice='imq'
+reg_weight=100
+#kernel_bandwidth=2
+# Hyperspherical VAE
+# Poincare VAE
+reconstruction_loss="bce"
+prior_distribution="riemannian_normal"
+posterior_distribution="wrapped_normal"
+curvature=0.7
+# Adversarial AE
+adversarial_loss_scale=0.9
+# VAE GAN
+#adversarial_loss_scale=0.8
+reconstruction_layer=3
+margin=0.4
+equilibrium= 0.68
+# VQ VAE
+commitment_loss_factor=0.25
+quantization_loss_factor=1.0
+num_embeddings=128
+use_ema=true
+decay=0.99
+# Hamiltonian VAE
+n_lf=1
+eps_lf=0.001
+beta_zero=0.3
+# Riemannian Hamiltonian VAE
+#n_lf=1
+#eps_lf=0.001
+#beta_zero=0.3
+temperature=1.5
+regularization=0.001
+# Importance Weighted Autoencoder
+number_samples=3
+# Multiply Importance Weighted Autoencoder
+number_gradient_estimates=4
+#number_samples=4
+# Partially Importance Weighted Autoencoder
+#number_gradient_estimates=4
+#number_samples=4
+# Combination Importance Weighted Autoencoder
+#beta=0.05
+#number_samples=4
+# VAMP Autoencoder
+number_components=50
+# Regularized AE with L2 decoder param
+embedding_weight=1e-2
+#reg_weight=1e-4
+# Regularized AE with gradient penalty
+#embedding_weight=1e-2
+#reg_weight=1e-4
+
 [Computational]
 gpu = true
 n_proc = 2