DeepSDFStruct 1.7.2__py3-none-any.whl

DeepSDFStruct/deep_sdf/workspace.py

@@ -0,0 +1,363 @@
+"""
+Experiment Workspace Management
+===============================
+
+This module provides utilities for managing DeepSDF experiment workspaces,
+including directory structures, file naming conventions, and model loading/saving.
+
+Constants
+---------
+The module defines standard subdirectory and file names for organizing
+experiment artifacts:
+- Model parameters and checkpoints
+- Optimizer states
+- Latent code vectors
+- Training logs and plots
+- Reconstructions and evaluations
+- Dataset samples and normalization parameters
+
+Architecture Registry
+--------------------
+ARCHITECTURES: dict
+    Maps architecture names to decoder classes, enabling dynamic model
+    instantiation from configuration files.
+
+Functions
+---------
+
+load_experiment_specifications
+    Load experiment configuration from specs.json file.
+
+load_trained_model
+    Load a trained decoder network from checkpoint.
+
+load_latent_vectors
+    Load learned latent codes from checkpoint.
+
+create_experiment_directory
+    Initialize directory structure for a new experiment.
+
+The workspace utilities ensure consistent organization across experiments
+and simplify model loading for inference and continued training.
+"""
+
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+import json
+import os
+import pathlib
+import torch
+from typing import TypedDict
+
+from .networks.analytic_round_cross import RoundCrossDecoder
+from .networks.deep_sdf_decoder import DeepSDFDecoder
+from .networks.hierarchical_deep_sdf_decoder import HierachicalDeepSDFDecoder
+from .networks.resnet_positional_sdf_decoder import ResNetPositionalDeepSDFDecoder
+from .networks.hierarchical_positional_sdf_decoder import (
+    HierachicalPositionalDeepSDFDecoder,
+)
+
+screenshots_subdir = "Screenshots"
+model_params_subdir = "ModelParameters"
+optimizer_params_subdir = "OptimizerParameters"
+latent_codes_subdir = "LatentCodes"
+latent_code_data_map_filename = "latent_code_data_map.json"
+logs_filename = "Logs.pth"
+logplot_filename = "Logs.png"
+reconstructions_subdir = "Reconstructions"
+reconstruction_meshes_subdir = "Meshes"
+reconstruction_codes_subdir = "Codes"
+specifications_filename = "specs.json"
+data_source_map_filename = ".datasources.json"
+evaluation_subdir = "Evaluation"
+sdf_samples_subdir = "SdfSamples"
+surface_samples_subdir = "SurfaceSamples"
+normalization_param_subdir = "NormalizationParameters"
+training_meshes_subdir = "TrainingMeshes"
+experiment_summary_name = "training_summary.json"
+
+# Map architecture name to Decoder class
+ARCHITECTURES = {
+    "analytic_round_cross": RoundCrossDecoder,
+    "deep_sdf_decoder": DeepSDFDecoder,
+    "hierarchical_deep_sdf_decoder": HierachicalDeepSDFDecoder,
+    "resnet_positional_deep_sdf_decoder": ResNetPositionalDeepSDFDecoder,
+    "hierarchical_positional_deep_sdf_decoder": HierachicalPositionalDeepSDFDecoder,
+}
+
+
+def load_experiment_specifications(experiment_directory):
+
+    filename = os.path.join(experiment_directory, specifications_filename)
+
+    if not os.path.isfile(filename):
+        raise Exception(
+            f"The experiment directory ({experiment_directory}) does not include specifications file "
+            + '"specs.json"'
+        )
+    with open(filename, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def load_latent_vectors(experiment_directory, checkpoint, device):
+
+    filename = os.path.join(
+        experiment_directory, latent_codes_subdir, checkpoint + ".pth"
+    )
+
+    if not os.path.isfile(filename):
+        raise Exception(
+            f"The experiment directory ({experiment_directory}) does not include a latent code file"
+            + f" for checkpoint '{checkpoint}'"
+        )
+
+    data = torch.load(filename, map_location=device, weights_only=True)
+
+    if isinstance(data["latent_codes"], torch.Tensor):
+
+        num_vecs = data["latent_codes"].size()[0]
+
+        lat_vecs = []
+        for i in range(num_vecs):
+            lat_vecs.append(data["latent_codes"][i])
+
+        return lat_vecs
+
+    else:
+
+        num_embeddings, embedding_dim = data["latent_codes"]["weight"].shape
+
+        lat_vecs = torch.nn.Embedding(num_embeddings, embedding_dim, device=device)
+
+        lat_vecs.load_state_dict(data["latent_codes"])
+
+        return lat_vecs.weight.data.detach()
+
+
+def load_model_parameters(
+    experiment_directory, checkpoint, decoder: torch.nn.Module, device
+):
+
+    filename = os.path.join(
+        experiment_directory, model_params_subdir, checkpoint + ".pth"
+    )
+
+    if not os.path.isfile(filename):
+        raise Exception('model state dict "{}" does not exist'.format(filename))
+
+    data = torch.load(filename, map_location=device, weights_only=True)
+
+    decoder.load_state_dict(data["model_state_dict"], strict=False)
+
+    return data["epoch"]
+
+
+def load_optimizer(
+    experiment_directory, checkpoint, optimizer: torch.nn.Module, device
+):
+
+    filename = os.path.join(
+        experiment_directory, optimizer_params_subdir, checkpoint + ".pth"
+    )
+
+    if not os.path.isfile(filename):
+        raise Exception(f'optimizer state dict "{filename}" does not exist')
+
+    data = torch.load(filename, map_location=device, weights_only=True)
+
+    optimizer.load_state_dict(data["optimizer_state_dict"])
+
+    return data["epoch"]
+
+
+def get_data_source_map_filename(data_dir):
+    return os.path.join(data_dir, data_source_map_filename)
+
+
+def get_reconstructed_mesh_filename(
+    experiment_dir,
+    epoch,
+    dataset,
+    class_name,
+    instance_name,
+    create_dir=True,
+    filetype="ply",
+):
+    fname_raw = os.path.join(
+        experiment_dir,
+        reconstructions_subdir,
+        str(epoch),
+        reconstruction_meshes_subdir,
+        dataset,
+        class_name,
+        instance_name + "." + filetype,
+    )
+    fname = pathlib.Path(fname_raw)
+    if not os.path.isdir(fname.parent) and create_dir:
+        os.makedirs(fname.parent)
+    return fname
+
+
+def get_reconstructed_code_filename(
+    experiment_dir, epoch, dataset, class_name, instance_name
+):
+
+    return os.path.join(
+        experiment_dir,
+        reconstructions_subdir,
+        str(epoch),
+        reconstruction_codes_subdir,
+        dataset,
+        class_name,
+        instance_name + ".pth",
+    )
+
+
+def get_evaluation_dir(experiment_dir, checkpoint, create_if_nonexistent=False):
+
+    dir = os.path.join(experiment_dir, evaluation_subdir, checkpoint)
+
+    if create_if_nonexistent and not os.path.isdir(dir):
+        os.makedirs(dir)
+
+    return dir
+
+
+def get_model_params_dir(experiment_dir, create_if_nonexistent=False):
+
+    dir = os.path.join(experiment_dir, model_params_subdir)
+
+    if create_if_nonexistent and not os.path.isdir(dir):
+        os.makedirs(dir)
+
+    return dir
+
+
+def get_screenshots_dir(experiment_dir, create_if_nonexistent=True):
+
+    dir = os.path.join(experiment_dir, screenshots_subdir)
+
+    if create_if_nonexistent and not os.path.isdir(dir):
+        os.makedirs(dir)
+
+    return dir
+
+
+def get_optimizer_params_dir(experiment_dir, create_if_nonexistent=False):
+
+    dir = os.path.join(experiment_dir, optimizer_params_subdir)
+
+    if create_if_nonexistent and not os.path.isdir(dir):
+        os.makedirs(dir)
+
+    return dir
+
+
+def get_latent_codes_dir(experiment_dir, create_if_nonexistent=False):
+
+    dir = os.path.join(experiment_dir, latent_codes_subdir)
+
+    if create_if_nonexistent and not os.path.isdir(dir):
+        os.makedirs(dir)
+
+    return dir
+
+
+def get_latent_code_data_map_filename(experiment_dir):
+    """Return absolute path for the latent-to-data mapping JSON file."""
+    return os.path.join(
+        get_latent_codes_dir(experiment_dir, create_if_nonexistent=True),
+        latent_code_data_map_filename,
+    )
+
+
+def get_normalization_params_filename(
+    data_dir, dataset_name, class_name, instance_name
+):
+    return os.path.join(
+        data_dir,
+        normalization_param_subdir,
+        dataset_name,
+        class_name,
+        instance_name + ".npz",
+    )
+
+
+def init_decoder(experiment_specs, device, data_parallel):
+    arch_name = experiment_specs["NetworkArch"]
+    if arch_name not in ARCHITECTURES:
+        raise ValueError(f"Unknown architecture: {arch_name}")
+
+    latent_size = experiment_specs["CodeLength"]
+    DecoderClass = ARCHITECTURES[arch_name]
+
+    decoder = DecoderClass(latent_size, **experiment_specs["NetworkSpecs"]).to(device)
+    if data_parallel:
+        decoder = torch.nn.DataParallel(decoder)
+    return decoder
+
+
+def load_trained_model(
+    experiment_directory: str, checkpoint: str, device=None, data_parallel=False
+):
+    specs_filename = os.path.join(experiment_directory, "specs.json")
+    with open(specs_filename, "r", encoding="utf-8") as f:
+        experiment_specs = json.load(f)
+    if device is None:
+        device = get_default_device(device)
+
+    filename = os.path.join(
+        experiment_directory, model_params_subdir, checkpoint + ".pth"
+    )
+
+    if not os.path.isfile(filename):
+        raise Exception('model state dict "{}" does not exist'.format(filename))
+
+    data = torch.load(filename, map_location=device)
+    decoder = init_decoder(experiment_specs, device, data_parallel)
+
+    state_dict = data["model_state_dict"]
+    if any(k.startswith("module.") for k in state_dict.keys()):
+        stripped_dict = {}
+        for k, v in state_dict.items():
+            new_key = k.replace("module.", "", 1)
+            stripped_dict[new_key] = v
+        state_dict = stripped_dict
+
+    decoder.load_state_dict(state_dict, strict=False)
+    decoder = decoder.to(device)
+    return decoder
+
+
+def get_default_device():
+
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    return device
+
+
+def print_model_specifications(experiment_directory: str):
+    specs = load_experiment_specifications(experiment_directory)
+    print("Model Specifications:")
+    for key in specs:
+        print(f"  {key}: {specs[key]}")
+    print("\n")
+
+
+class ExperimentSummary(TypedDict):
+    loss: float
+    num_epochs: int
+    timestamp: str
+    host_name: str
+    device: str
+    training_duration: str
+    data_dir: str
+    version: str
+
+
+def save_experiment_summary(experiment_directory: str, summary: ExperimentSummary):
+    with open(os.path.join(experiment_directory, experiment_summary_name), "w") as f:
+        json.dump(summary, f, indent=4)

DeepSDFStruct/design_of_experiments.py

@@ -0,0 +1,235 @@
+"""
+Design of Experiments (DOE) for DeepSDF Training
+================================================
+
+This module provides tools for conducting systematic design of experiments
+when training DeepSDF neural network models. It integrates with MLflow for
+experiment tracking and supports automated hyperparameter sweeps.
+
+Key Features
+------------
+
+ExperimentSpecifications
+    A flexible dictionary-like class for managing experiment configurations
+    with support for:
+    - Loading specifications from JSON files
+    - Recursive updates for nested parameter structures
+    - Deep copying to prevent unintended modifications
+    - Flattening nested configurations for logging
+    - Saving configurations for reproducibility
+
+MLflow Integration
+    Automatic tracking of:
+    - Training metrics (loss curves, validation scores)
+    - Model checkpoints
+    - Hyperparameters and configurations
+    - Dataset information
+    - Experiment metadata
+
+The module simplifies the process of running large-scale hyperparameter
+searches and ablation studies for DeepSDF model training, with built-in
+versioning and reproducibility features.
+
+Examples
+--------
+Run an experiment with custom specifications::
+
+    from DeepSDFStruct.design_of_experiments import ExperimentSpecifications
+
+    # Load base configuration
+    specs = ExperimentSpecifications('config.json')
+
+    # Update hyperparameters
+    specs.update({
+        'learning_rate': 0.001,
+        'batch_size': 32,
+        'num_epochs': 100
+    })
+
+    # Run training with tracking
+    # train_with_specs(specs)
+
+Create and modify experiment configurations::
+
+    specs = ExperimentSpecifications({
+        'model': {
+            'layers': [512, 512, 512],
+            'activation': 'relu'
+        },
+        'training': {
+            'lr': 0.0005
+        }
+    })
+
+    # Flatten for logging
+    flat_params = specs.flatten()
+    print(flat_params)  # {'model.layers': [...], 'model.activation': 'relu', ...}
+"""
+
+import os
+import pathlib
+import json
+import mlflow
+from urllib.parse import urlparse
+from DeepSDFStruct.deep_sdf.training import train_deep_sdf
+import copy
+from typing import Dict, Any, Union
+
+package_path = pathlib.Path(__file__).parent
+
+
+class ExperimentSpecifications(dict):
+    """
+    A dictionary-like class to hold and update experiment specifications.
+    Supports recursive updates for nested dictionaries.
+    Can be initialized from a dictionary or loaded from a file (JSON/YAML).
+    """
+
+    def __init__(self, specs: Union[Dict[str, Any], str, None] = None):
+        """
+        Initialize the ExperimentSpecifications.
+
+        Args:
+            specs: A dictionary of specifications, a file path (JSON/YAML),
+                   or None for an empty specification set.
+        """
+        if isinstance(specs, str):
+            loaded_specs = self._load_from_file(specs)
+            super().__init__(copy.deepcopy(loaded_specs))
+        else:
+            super().__init__(copy.deepcopy(specs) if specs else {})
+
+    @staticmethod
+    def _load_from_file(filename: str) -> Dict[str, Any]:
+        """
+        Load specifications from a JSON or YAML file.
+        """
+        with open(filename, "r") as f:
+            if filename.endswith(".json"):
+                return json.load(f)
+            else:
+                raise ValueError("Unsupported file format. Use .json or .yaml/.yml")
+
+    def save(self, filename: str) -> None:
+        """
+        Save current specifications to a JSON or YAML file.
+        """
+        with open(filename, "w") as f:
+            if filename.endswith(".json"):
+                json.dump(self, f, indent=4)
+
+    def update(self, updates: Dict[str, Any]) -> None:
+        """
+        Recursively update the experiment specifications with new values.
+        """
+
+        def _recursive_update(d, u):
+            for k, v in u.items():
+                if isinstance(v, dict) and isinstance(d.get(k), dict):
+                    _recursive_update(d[k], v)
+                else:
+                    d[k] = v
+
+        _recursive_update(self, updates)
+
+    def copy(self):
+        """Return a deep copy of the specifications."""
+        return ExperimentSpecifications(copy.deepcopy(self))
+
+    def flatten(self, parent_key: str = "", sep: str = ".") -> Dict[str, Any]:
+        """
+        Flatten nested dictionary (including lists of dicts) into a single-level dict
+        with dot-separated keys.
+        """
+        items = {}
+        for k, v in self.items():
+            new_key = f"{parent_key}{sep}{k}" if parent_key else k
+
+            if isinstance(v, dict):
+                items.update(
+                    ExperimentSpecifications(v).flatten(parent_key=new_key, sep=sep)
+                )
+
+            elif isinstance(v, list):
+                # Check if the list contains dictionaries
+                if all(isinstance(i, dict) for i in v):
+                    for idx, elem in enumerate(v):
+                        items.update(
+                            ExperimentSpecifications(elem).flatten(
+                                parent_key=f"{new_key}.{idx}", sep=sep
+                            )
+                        )
+                else:
+                    # Just convert the list/tuple to a string for MLflow
+                    items[new_key] = str(v)
+
+            else:
+                items[new_key] = v
+
+        return items
+
+    def __repr__(self):
+        return f"ExperimentSpecifications({dict(self)})"
+
+
+def create_experiment(exp_dir, specs):
+    """
+    Create a new experiment directory and specs.json file, copying defaults and applying overrides.
+
+    Args:
+        exp_dir (str): Path to new experiment directory.
+        specs (dict): Dictionary containing the experiment specifications.
+    """
+
+    os.makedirs(exp_dir, exist_ok=True)
+
+    # Write new specs.json
+    specs_path = os.path.join(exp_dir, "specs.json")
+    with open(specs_path, "w") as f:
+        json.dump(specs, f, indent=4)
+
+    return specs_path
+
+
+def run_experiment(
+    exp_name,
+    data_dir,
+    specs: ExperimentSpecifications,
+    batch_split=1,
+    run_name=None,
+    device="cpu",
+    tracking_uri="sqlite:///mlruns.db",
+):
+    """
+    Creates an experiment with overrides, trains the model, and logs everything to MLflow.
+
+    Args:
+        exp_name (str): Name of the experiment folder.
+        data_dir (str): Path to the dataset.
+        specs (dict): Experiment specifications
+        device (str): Training device ('cpu' or 'cuda').
+        mlflow_experiment_name (str): Name of the MLflow experiment.
+        register_model (bool): If True, register the final model in MLflow Model Registry.
+    """
+    mlflow.set_tracking_uri(tracking_uri)
+    # Start MLflow run
+    mlflow.set_experiment(exp_name)
+    with mlflow.start_run(run_name=run_name):
+        parsed = urlparse(mlflow.get_artifact_uri())
+        if parsed.path == "":
+            raise NotImplementedError("Remote Tracking URI not implemented yet.")
+        else:
+            exp_dir = parsed.path
+        create_experiment(exp_dir, specs=specs)
+        specs_path = os.path.join(exp_dir, "specs.json")
+        mlflow.log_artifact(specs_path)
+
+        mlflow.log_params(specs.flatten())
+
+        summary = train_deep_sdf(
+            exp_dir, data_dir, device=device, batch_split=batch_split
+        )
+        mlflow.set_tags(dict(summary))
+
+        mlflow.log_metric("train_loss", summary["loss"])
+    return summary["loss"]

DeepSDFStruct/flexicubes/__init__.py

@@ -0,0 +1,15 @@
+"""
+FlexiCubes - Differentiable 3D Mesh Extraction
+==============================================
+
+This module implements FlexiCubes, a differentiable 3D variant of the
+Dual Marching Cubes algorithm for extracting meshes from scalar fields.
+
+FlexiCubes enables gradient-based optimization of mesh representations
+by making the mesh extraction process fully differentiable. This is
+essential for inverse problems and optimization tasks where the geometry
+must be optimized to satisfy certain constraints or objectives.
+
+The module uses precomputed lookup tables to handle all 256 possible
+Marching Cubes configurations efficiently.
+"""