manifold-microscope 0.0.1__py3-none-any.whl

experiment_scripts/manifold_analysis/analysis.py

@@ -0,0 +1,107 @@
+import itertools
+from pathlib import Path
+
+import typer
+from tqdm import tqdm
+
+from microscope.datasets.generic_dataset_loader import DatasetName
+from microscope.computations_grid.data_analysis.run_data_analysis import main as run_analysis
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+def analysis_on_model(
+    inference_path: Path,
+    output_path: Path,
+    dataset_name: DatasetName,
+    model_type: str,
+    number_of_dims: int,
+    only_evolution: bool,
+    normalize_for_volume: bool,
+    skip_done: bool,
+    n_samples_for_plots: int = 50_000
+) -> None:
+    if output_path.exists() and skip_done:
+        print(f"Skipping {output_path.name} as it exists.")
+        return None
+
+    print(f"Analysis on {output_path.name}.")
+    run_analysis(
+        inference_path=inference_path,
+        output_path=output_path,
+        dataset=dataset_name,
+        model_type=model_type,
+        number_of_dims=number_of_dims,
+        only_evolution=only_evolution,
+        normalize_for_volume=normalize_for_volume,
+        n_samples_for_plots=n_samples_for_plots
+    )
+
+
+@app.command()
+def run_analyses(
+    inference_path: Path = typer.Option(...),
+    output_path: Path = typer.Option(...),
+    only_evolution: bool = True,
+    normalize_for_volume: bool = True,
+    skip_done: bool = True,
+    n_samples_for_plots: int = 50_000
+) -> None:
+    dataset_name_list = [
+        "custom_dsprites_balanced",
+        "extended_coil20"
+    ]
+    model_type_list = [
+        "beta_vae",
+        "mae"
+    ]
+    training_ratio_per_dim_list = [
+        1.0
+    ]
+    number_of_dims_list = [
+        # 1,
+        2,
+        3,
+        # 4
+    ]
+    # No noise.
+    noise_sigma_list = [
+        0
+    ]
+
+    hyperparameter_grid = list(itertools.product(
+        dataset_name_list,
+        model_type_list,
+        training_ratio_per_dim_list,
+        number_of_dims_list,
+        noise_sigma_list
+    ))
+
+    for dataset_name, model_type, training_ratio_per_dim, number_of_dims, noise_sigma in tqdm(hyperparameter_grid):
+        # Skip dimension 4 for COIL20.
+        if (number_of_dims == 4) and (dataset_name == "extended_coil20"):
+            continue
+        if (number_of_dims == 4) and (model_type == "mae"):
+            continue
+        model_dir = "__".join([
+            dataset_name,
+            model_type,
+            str(training_ratio_per_dim),
+            str(number_of_dims),
+            str(noise_sigma)
+        ])
+        analysis_on_model(
+            output_path=output_path / model_dir,
+            inference_path=inference_path / model_dir,
+            dataset_name=dataset_name,
+            model_type=model_type,
+            number_of_dims=number_of_dims,
+            only_evolution=only_evolution,
+            normalize_for_volume=normalize_for_volume,
+            skip_done=skip_done,
+            n_samples_for_plots=n_samples_for_plots
+        )
+
+
+if __name__ == "__main__":
+    app()

experiment_scripts/manifold_fitting/analysis.py

@@ -0,0 +1,106 @@
+import itertools
+from pathlib import Path
+
+import typer
+from tqdm import tqdm
+
+from microscope.datasets.generic_dataset_loader import DatasetName
+from microscope.computations_grid.data_analysis.run_data_analysis import main as run_analysis
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+def analysis_on_model(
+    inference_path: Path,
+    output_path: Path,
+    dataset_name: DatasetName,
+    model_type: str,
+    number_of_dims: int,
+    only_evolution: bool,
+    normalize_for_volume: bool,
+    skip_done: bool
+) -> None:
+    if output_path.exists() and skip_done:
+        print(f"Skipping {output_path.name} as it exists.")
+        return None
+
+    print(f"Analysis on {output_path.name}.")
+    run_analysis(
+        inference_path=inference_path,
+        output_path=output_path,
+        dataset=dataset_name,
+        model_type=model_type,
+        number_of_dims=number_of_dims,
+        only_evolution=only_evolution,
+        normalize_for_volume=normalize_for_volume
+    )
+
+
+@app.command()
+def run_analyses(
+    inference_path: Path = typer.Option(...),
+    output_path: Path = typer.Option(...),
+    only_evolution: bool = True,
+    normalize_for_volume: bool = True,
+    skip_done: bool = True
+) -> None:
+    dataset_name_list = [
+        "custom_dsprites_balanced",
+        "extended_coil20"
+    ]
+    model_type_list = [
+        "beta_vae",
+        "mae"
+    ]
+    training_ratio_per_dim_list = [
+        0.4,
+        0.5,
+        0.6,
+        1.0
+    ]
+    number_of_dims_list = [
+        1,
+        2,
+        3,
+        4
+    ]
+    # No noise.
+    noise_sigma_list = [
+        0
+    ]
+
+    hyperparameter_grid = list(itertools.product(
+        dataset_name_list,
+        model_type_list,
+        training_ratio_per_dim_list,
+        number_of_dims_list,
+        noise_sigma_list
+    ))
+
+    for dataset_name, model_type, training_ratio_per_dim, number_of_dims, noise_sigma in tqdm(hyperparameter_grid):
+        # Skip dimension 4 for COIL20.
+        if (number_of_dims == 4) and (dataset_name == "extended_coil20"):
+            continue
+        if (number_of_dims == 4) and (model_type == "mae"):
+            continue
+        model_dir = "__".join([
+            dataset_name,
+            model_type,
+            str(training_ratio_per_dim),
+            str(number_of_dims),
+            str(noise_sigma)
+        ])
+        analysis_on_model(
+            output_path=output_path / model_dir,
+            inference_path=inference_path / model_dir,
+            dataset_name=dataset_name,
+            model_type=model_type,
+            number_of_dims=number_of_dims,
+            only_evolution=only_evolution,
+            normalize_for_volume=normalize_for_volume,
+            skip_done=skip_done
+        )
+
+
+if __name__ == "__main__":
+    app()

experiment_scripts/manifold_fitting/inference.py

@@ -0,0 +1,147 @@
+import itertools
+from pathlib import Path
+
+import typer
+from tqdm import tqdm
+
+from microscope.datasets.generic_dataset_loader import DatasetName
+from representation_learning.mae.inference_intermediate_layers import main as inference_mae
+from representation_learning.beta_vae.inference_intermediate_layers import main as inference_beta_vae
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+def inference_on_model(
+    output_path: Path,
+    model_path: Path,
+    only_final_model: bool,
+    only_latent_and_output: bool,
+    dataset_name: DatasetName,
+    model_type: str,
+    number_of_dims: int,
+    skip_done: bool
+) -> None:
+    if output_path.exists() and skip_done:
+        print(f"Skipping {output_path.name} as it exists.")
+        return None
+
+    print(f"Inference on {output_path.name}.")
+    if model_type == "beta_vae":
+        inference_fn = inference_beta_vae
+    elif model_type == "mae":
+        inference_fn = inference_mae
+    else:
+        raise ValueError(f"Unknown model type: {model_type}.")
+
+    if only_final_model:
+        checkpoints_path = model_path / "checkpoints"
+        final_candidates = [
+            ckpt for ckpt in checkpoints_path.glob("*")
+            if "last" in ckpt.name or "final" in ckpt.name
+        ]
+        if len(final_candidates) != 1:
+            raise ValueError(f"Found the following final model candidates {final_candidates}. Expected as single one.")
+        checkpoint_path = final_candidates[0]
+        random_model = False
+
+        inference_fn(
+            dataset=dataset_name,
+            number_of_dims=number_of_dims,
+            only_latent_and_output=only_latent_and_output,
+            checkpoint_path=checkpoint_path,
+            output_dir=output_path,
+            random_model=random_model
+        )
+    else:
+        checkpoints_path = model_path / "checkpoints"
+        checkpoint_paths = [p for p in checkpoints_path.glob("*") if "npz" not in p.suffix]
+        final_candidates = [
+            ckpt for ckpt in checkpoint_paths
+            if "last" in ckpt.name or "final" in ckpt.name
+        ]
+        if len(final_candidates) != 1:
+            raise ValueError(f"Found the following final model candidates {final_candidates}. Expected as single one.")
+        final_checkpoint = final_candidates[0]
+        middle_checkpoint_idx = len(checkpoint_paths) // 2
+        middle_checkpoint = checkpoint_paths[middle_checkpoint_idx]
+        first_checkpoint = checkpoint_paths[0]
+
+        for checkpoint_path in [first_checkpoint, middle_checkpoint, final_checkpoint]:
+            inference_fn(
+                dataset=dataset_name,
+                number_of_dims=number_of_dims,
+                only_latent_and_output=only_latent_and_output,
+                checkpoint_path=checkpoint_path,
+                output_dir=output_path,
+                random_model=False
+            )
+
+
+@app.command()
+def run_inferences(
+    training_path: Path = typer.Option(...),
+    output_path: Path = typer.Option(...),
+    only_final_model: bool = True,
+    only_latent_and_output: bool = True,
+    skip_done: bool = True
+) -> None:
+    dataset_name_list = [
+        "custom_dsprites_balanced",
+        "extended_coil20"
+    ]
+    model_type_list = [
+        "beta_vae",
+        "mae"
+    ]
+    training_ratio_per_dim_list = [
+        0.4,
+        0.5,
+        0.6,
+        1.0
+    ]
+    number_of_dims_list = [
+        1,
+        2,
+        3,
+        4
+    ]
+    # No noise.
+    noise_sigma_list = [
+        0
+    ]
+
+    hyperparameter_grid = list(itertools.product(
+        dataset_name_list,
+        model_type_list,
+        training_ratio_per_dim_list,
+        number_of_dims_list,
+        noise_sigma_list
+    ))
+
+    for dataset_name, model_type, training_ratio_per_dim, number_of_dims, noise_sigma in tqdm(hyperparameter_grid):
+        # Skip dimension 4 for COIL20.
+        if (number_of_dims == 4) and (dataset_name == "extended_coil20"):
+            continue
+        if (number_of_dims == 4) and (model_type == "mae"):
+            continue
+        model_dir = "__".join([
+            dataset_name,
+            model_type,
+            str(training_ratio_per_dim),
+            str(number_of_dims),
+            str(noise_sigma)
+        ])
+        inference_on_model(
+            output_path=output_path / model_dir,
+            model_path=training_path / model_dir,
+            only_final_model=only_final_model,
+            only_latent_and_output=only_latent_and_output,
+            dataset_name=dataset_name,
+            model_type=model_type,
+            number_of_dims=number_of_dims,
+            skip_done=skip_done
+        )
+
+
+if __name__ == "__main__":
+    app()

experiment_scripts/manifold_fitting/mmls.py

@@ -0,0 +1,63 @@
+import pickle
+from pathlib import Path
+
+import numpy as np
+
+from experiment_scripts.model_configs import MMLSConfig
+from experiment_scripts.toy_manifolds_experiment.manifold_fitting_no_noise import ANNMMLSProjector
+from microscope.datasets.generic_dataset_loader import DatasetName, load_dataset_fixed_test_split
+
+
+def fit_mmls(model_config: MMLSConfig) -> None:
+    output_dir = Path(model_config.output_dir)
+    exported_datasets_dir = Path(model_config.exported_datasets_dir)
+    dataset = DatasetName[model_config.dataset]
+    number_of_dims = model_config.number_of_dims
+    training_ratio = model_config.training_ratio
+    ratio_per_dim = model_config.ratio_per_dim
+    noise_sigma = model_config.noise_sigma
+    number_of_neighbors = model_config.number_of_neighbors
+    verbose = model_config.verbose
+    device = model_config.device
+
+    config = locals()
+
+    # Load the exported dataset.
+    data_train, data_test, _, _ = load_dataset_fixed_test_split(
+        datasets_dir=exported_datasets_dir,
+        dataset_name=dataset,
+        number_of_dims=number_of_dims,
+        ratio_per_dim=ratio_per_dim,
+        training_ratio=training_ratio,
+        noise_sigma=noise_sigma,
+        weight_subsampling_by_manifold_volume=True
+    )
+
+    # Load the dataset.
+    data_train = (data_train - data_train.min()) / (data_train.max() - data_train.min())
+    data_test = (data_test - data_test.min()) / (data_test.max() - data_test.min())
+    data_train = data_train.reshape(data_train.shape[0], np.prod(data_train.shape[1:]))
+    data_test = data_test.reshape(data_test.shape[0], np.prod(data_test.shape[1:]))
+
+    if len(data_train) < number_of_neighbors:
+        print(
+            f"Skipping the training ratio {training_ratio} as it results to {len(data_train)} points which are less "
+            f"than the number of neighbors {number_of_neighbors}."
+        )
+        return
+
+    projector = ANNMMLSProjector(data_train, d=number_of_dims, k=number_of_neighbors, verbose=verbose, device=device)
+    prediction = projector.project(data_test)
+
+    distances = np.linalg.norm(data_test - prediction, axis=-1)
+    hausdorff_distance = distances.max()
+
+    results = dict(
+        number_of_train_points=len(data_train),
+        number_of_test_points=len(data_test),
+        pointwise_distances=distances,
+        hausdorff_distance=hausdorff_distance
+    )
+
+    with open(output_dir / f"distance_results_{dataset}.pkl", "wb") as f:
+        pickle.dump(results, f, -1)

experiment_scripts/manifold_fitting/training.py

@@ -0,0 +1,220 @@
+import itertools
+from dataclasses import asdict, replace
+from pathlib import Path
+
+import numpy as np
+import torch
+import typer
+import yaml
+
+from experiment_scripts.manifold_fitting.mmls import fit_mmls
+from experiment_scripts.model_configs import BetaVAEConfig, MMLSConfig
+from microscope.datasets.generic_dataset_loader import DatasetName, export_fixed_grid_test_set_and_rest_for_train
+from representation_learning.beta_vae.solver import Solver
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+def update_max_epochs(
+    config: BetaVAEConfig,
+    mini_test_run: bool,
+    training_ratio: float
+) -> BetaVAEConfig:
+    if mini_test_run:
+        config = replace(config, max_epochs=1, plot_interval=2)
+    elif training_ratio < 1.0:
+        ratio_correction_coeff = 1 / training_ratio
+        max_epochs = int(ratio_correction_coeff * config.max_epochs)
+        plot_interval = int(ratio_correction_coeff * config.plot_interval)
+        config = replace(
+            config,
+            max_epochs=max_epochs,
+            plot_interval=plot_interval
+        )
+
+    return config
+
+
+def train_model(
+    output_path: Path,
+    exported_datasets_dir: Path,
+    seed: int,
+    dataset_name: str,
+    model_type: str,
+    training_ratio: float,
+    ratio_per_dim: bool,
+    number_of_dims: int,
+    noise_sigma: float,
+    skip_done: bool,
+    device: str = "cpu"
+) -> None:
+    if output_path.exists() and skip_done:
+        print(f"Skipping {output_path.name} as it exists.")
+        return None
+
+    print(f"Training on {output_path.name}.")
+    output_path.mkdir(parents=True)
+    experiment_config = dict(
+        output_path=str(output_path),
+        seed=seed,
+        dataset_name=dataset_name,
+        model_type=model_type,
+        training_ratio=training_ratio,
+        ratio_per_dim=ratio_per_dim,
+        number_of_dims=number_of_dims,
+        noise_sigma=noise_sigma,
+        device=device,
+    )
+
+    match model_type:
+        case "beta_vae":
+            model_config = BetaVAEConfig(
+                dataset=dataset_name,
+                ckpt_dir=output_path / "checkpoints",
+                exported_datasets_dir=exported_datasets_dir,
+                output_dir=output_path,
+                training_ratio=training_ratio,
+                ratio_per_dim=ratio_per_dim,
+                number_of_dims=number_of_dims,
+                noise_sigma=noise_sigma,
+                device=device
+            )
+
+            if dataset_name == DatasetName.extended_coil20:
+                model_config = replace(
+                    model_config,
+                    max_epochs=int(1e5),
+                    objective="H",
+                    model="H",
+                    lr=1e-4,
+                    loss_threshold=10
+                )
+
+            # Save config.
+            config = dict(
+                experiment_config=experiment_config,
+                model_config=asdict(model_config)
+            )
+            with open(output_path / "config.yml", "w") as f:
+                yaml.dump(config, f)
+
+            torch.manual_seed(seed)
+            if device == "mps" and torch.backends.mps.is_available():
+                torch.mps.manual_seed(seed)
+            elif device.startswith("cuda") and torch.cuda.is_available():
+                torch.cuda.manual_seed(seed)
+            np.random.seed(seed)
+
+            net = Solver(args=model_config)
+            net.train()
+        case "MMLS":
+            model_config = MMLSConfig(
+                output_dir=str(output_path),
+                exported_datasets_dir=exported_datasets_dir,
+                dataset=dataset_name,
+                number_of_dims=number_of_dims,
+                training_ratio=training_ratio,
+                ratio_per_dim=ratio_per_dim,
+                noise_sigma=noise_sigma,
+                number_of_neighbors=2*2**number_of_dims,
+                device=device
+            )
+            fit_mmls(model_config)
+        case _:
+            raise ValueError(f"Unknown model type: {model_type}.")
+
+
+@app.command()
+def run_trainings(
+    output_path: Path = typer.Option(...),
+    n_experiment_repetitions: int = 1,
+    max_test_size: int = 500,
+    seed: int = 42,
+    skip_done: bool = True,
+    device: str = "cuda:0"
+) -> None:
+    exported_datasets_dir = output_path / "datasets"
+    exported_datasets_dir.mkdir(exist_ok=True, parents=True)
+    dataset_name_list = [
+        "custom_dsprites_balanced",
+        "extended_coil20"
+    ]
+    model_type_list = [
+        "beta_vae",
+        "MMLS"
+    ]
+    ratio_per_dim = False
+    # No noise.
+    noise_sigma_list = [
+        0
+    ]
+
+    hyperparameter_grid = list(itertools.product(
+        dataset_name_list,
+        model_type_list,
+        noise_sigma_list
+    ))
+
+    for dataset_name, model_type, noise_sigma in hyperparameter_grid:
+        match dataset_name:
+            case "extended_coil20":
+                training_ratio_list = [0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+
+                number_of_dims_list = [
+                    1,
+                    2,
+                    3
+                ]
+            case "custom_dsprites_balanced":
+                training_ratio_list = [0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+                number_of_dims_list = [
+                    1,
+                    2,
+                    3,
+                    4
+                ]
+            case _:
+                raise ValueError(f"Unknown dataset name: {dataset_name}.")
+
+        second_level_hyperparameter_grid = list(itertools.product(
+            number_of_dims_list,
+            training_ratio_list
+        ))
+        for number_of_dims, training_ratio in second_level_hyperparameter_grid:
+            export_fixed_grid_test_set_and_rest_for_train(
+                dataset_name=dataset_name,
+                number_of_dims=number_of_dims,
+                output_dir=exported_datasets_dir,
+                max_test_size=max_test_size,
+                device=device
+            )
+
+            for repetition_n in range(n_experiment_repetitions):
+                if number_of_dims == 1:
+                    if training_ratio < 0.1:
+                        continue
+                model_dir = "__".join([
+                    dataset_name,
+                    model_type,
+                    str(round(training_ratio, 3)),
+                    str(number_of_dims),
+                    str(noise_sigma),
+                    str(repetition_n)
+                ])
+                train_model(
+                    output_path=output_path / model_dir,
+                    exported_datasets_dir=exported_datasets_dir,
+                    seed=seed,
+                    dataset_name=dataset_name,
+                    model_type=model_type,
+                    training_ratio=training_ratio,
+                    ratio_per_dim=ratio_per_dim,
+                    number_of_dims=number_of_dims,
+                    noise_sigma=noise_sigma,
+                    skip_done=skip_done,
+                    device=device
+                )
+
+
+if __name__ == "__main__":
+    app()