Unifi3D: A Study on 3D Representations for Generation and Reconstruction in a Common Framework

This is the code repository for Unifi3D, a unified framework for evaluating the reconstruction and generation performance of 3D representations. We compare these representations based on multiple criteria: quality, computational efficiency, and generalization performance. Beyond standard model benchmarking, our experiments aim to derive best practices over all steps involved in the 3D generation pipeline, including preprocessing, mesh reconstruction, compression with autoencoders, and generation. Our findings highlight that reconstruction errors significantly impact overall performance, underscoring the need to evaluate generation and reconstruction jointly.

1.1: Repository structure

Details (Click to expand)

The Unifi3D repository has the following structure:

Unifi3D                                   # The top-level Unifi3D repository.
|                                         #
├── unifi3d                               # Core module regrouping the main functionalities of Unifi3D
|    |                                    #
|    ├── data                             # Torch dataset & sampler classes
|    |
|    ├── losses                           # Loss functions including method specific functions.
|    |   |
|    │   ├── diffusion                    #
|    |   └── autoencoders                 #
|    |
|    ├── models                           # Relevant neural models and related submodules
|    │   │                                #
|    │   ├── diffusion                    #
|    │   ├── autoencoders                 #
|    |   └── modules                      # Representation, personalized network layers, or model componeents
|    |
|    ├── utils                            # Useful functions
|    │   │                                #
|    │   ├── data                         #
|    │   ├── evaluate                     #
|    │   ├── logger                       #
|    │   ├── rendering                    #
|    │   ├── model                        #
|    │   ├── triplane_utils               #
|    |   └── scheduler                    #
|
├── configs                               # Hydra configs
|    |                                    #
|    ├── data                             # Data configs
|    ├── debug                            # Debugging configs
|    ├── experiment                       # Experiment configs
|    ├── extras                           # Extra utilities configs
|    ├── hparams_search                   # Hyperparameter search configs
|    ├── hydra                            # Hydra configs
|    ├── local                            # Local configs
|    ├── logger                           # Logger configs
|    ├── model                            # Model configs
|    ├── paths                            # Project paths configs
|    ├── trainer                          # Trainer configs
|    │                                    #
|    ├── default.yaml                     # Config template
|    ├── eval.yaml                        # Main config for evaluation
|    ├── benchmark.yaml                   # Main config for benchmarking reconstruction
|    ├── sample_diffusion.yaml            # Main config for generation (inference)
|    └── train.yaml                       # Main config for training
|                                         #
├── checkpoints                           # Symbolic links to saved checkpoints in the format of *.safetensors. 
|                                         #
├── data                                  # Torch dataset & sampler classes
│   ├── split_shapenet.csv                # our train-test split of the ShapeNet dataset
│   ├── ShapeNetCore.v1                   # original dataset (NOTE: to be downloaded)
│   └── shapenet_preprocessed             # preprocessed dataset (NOTE: to be created)
|
├── docs                                  # Additional documentation markdown files
|                                         #
├── logs                                  # Result of the training and evaluation runs
|                                         #
├── media                                 # Pictures and other relevant media resources of this repository
|                                         #
├── scripts                               # Python scripts and Jupyter notebooks
|    |                                    #
|    ├── accelerate                       # Main training scripts
|    ├── data                             # Dataset preprocessing scripts
|    ├── evaluate                         # Scripts for evaluation and creating tables and plots
|    ├── diffusion                        # Diffusion inference script
|    └── utils                            # Useful general scripts
│                                         #
├── tests                                 # Unit tests
│                                         #
├── .project-root                         # File for inferring the position of project root directory
├── .python-version                       # Contains the exact version of python used in the project
│                                         #
├── requirements_<platform>.txt           # File containing some of the main dependencies of the project by platform
├── LICENSE                               # License file
├── README.md                             # ReadMe file
└── setup.py                              # File for installing project as a package

1.2: Installation & Dependencies

Details (Click to expand)

We provide a bash script allowing users to automatically setup unifi3d and to create a suitable virtual environment for the project, connected to a specific and well tested version of python. This is the recommended way. Execute the script with:

chmod a+x setup_unifi3d.sh
./setup_unifi3d.sh

Note: Some components, such as dataset creation and user study scripts, require additional dependencies not installed by the setup script:

• Blender 4.1 (ensure it is available in your PATH)

For alternative installation methods (using venv, manual step-by-step setup, etc.), refer to our setup guide.

2: Dataset

Details (Click to expand)

We use the ShapeNet dataset for our experiments. To enable different 3D generation approaches, we preprocess the data into point cloud, SDF etc. Details on generated datasets.

To preprocess the data, follow these steps:

• Download the ShapeNetCore.v1 dataset here.
• Once downloaded, place it in the data directory (there should be a folder data/ShapeNetCore.v1).
• Install blender (if not already done) and add it to the PATH (see setup guide).
• Run python scripts/data/create_preprocessed_dataset.py --task_id 0 --num_tasks 1 data/shapenet_preprocessed (Note: This can take a while. We recommend to distribute the work to multiple parallel jobs using a large number of CPU cores or multiple machines. Use the --task_id and --num_tasks to split the work for your parallel compute environment.)

This will save the preprocessed data in data/shapenet_preprocessed.

3: Training pipeline

Details (Click to expand)

• Start by activating the unifi3d virtual environment:

  conda activate unifi3d

• Setup an experiment parameter .yaml file in configs/experiment, for instance my_great_experiment.yaml. Within this file, set the parameters of your experiment.

• Run the trainer (used for both AE and diffusion training)

  python scripts/accelerate/acc_train.py experiment=my_great_experiment

  The training script should start automatically, instanciate a Trainer and a Logger based on the specifications of your experiment parameter file my_great_experiment.yaml.

• In case you are still in debug phase and with to see the crash logs, you may use the HYDRA_FULL_ERROR=1 before your command::

  HYDRA_FULL_ERROR=1 python scripts/accelerate/acc_train.py experiment=my_great_experiment

• For multi-GPU training, we rely on the hugging-face accelerate pipeline.

  • You may start setting up your accelerate environment by creating an accelerate configuration file:

    accelerate config

    This will generate a default_config.yaml file in $HOME/.cache/huggingface/accelerate which will be read by accelerate at the beginning of your training session. Simply follow the setup guide and set the number of GPUs according to your hardware.
  • Then on your multi-GPU mahcine or cluster session, you may run the distributed experiment using the following command:

    accelerate launch scripts/accelerate/acc_train.py experiment=my_great_experiment

  • Note that you can point to a specific config file to change the parameters of accelerate (e.g. for training on a different number of GPUs) using the --config_file flag:

    accelerate launch --config_file path/to/config/my_config_file.yaml scripts/accelerate/acc_train.py experiment=my_great_experiment

  • You might also want to avoid using the accelerate config file for greater flexibility. In this case, you may simply use the --multi_gpu flag:

    accelerate launch --multi_gpu scripts/accelerate/acc_train.py experiment=my_great_experiment

    In this case, Accelerate will make some hyperparameter decisions for you, e.g., if GPUs are available, it will use all of them by default without the mixed precision.
  • For a complete list of parameters you can pass in, run:

    accelerate launch -h

    Check the hugging-face documentation for more details.

Currently, all config files for AE model training (i.e., using EncoderDecoderTrainer) have the acc_ prefix, e.g. acc_shape2vecset_ae, while all diffusion models have the _diffuse prefix, e.g. diffuse_sdf_ae_dit.

Most models were trained on 48 GB GPUs (e.g. single GPU with batch size 8 for SDF, Voxel, Triplane & Shap2vecset. DualOctree can be trained with larger batch sizes, Shap-E requires 80GB).

3.1 Training encoder models

Example usage:

python scripts/accelerate/acc_train.py experiment=acc_triplane_ae
python scripts/accelerate/acc_train.py experiment=acc_shape2vecset_ae
python scripts/accelerate/acc_train.py experiment=acc_sdf_ae
python scripts/accelerate/acc_train.py experiment=acc_dualoctree_vae
python scripts/accelerate/acc_train.py experiment=acc_shap_e_ae

3.2 Training diffusion models

When training a diffusion model, you need to specify the checkpoint of the pretrained autoencoder. For example, train a DiT to denoise an encoded SDF with

python scripts/accelerate/acc_train.py experiment=diffuse_sdf_dit net_encode.ckpt_path=/path/to/ckpt/model.safetensors

Or for DualOctree with VAE and Unet:

python scripts/accelerate/acc_train.py experiment=diffuse_doctree_vae_unet net_encode.ckpt_path=/path/to/ckpt/model.safetensors

3.3 Visualize aim stack logging

• Enable tunneling with your desired terminal and set the destination server to 127.0.0.1:43800:

  cd ~/unifi3d/
  aim up
  

4: Benchmarking reconstruction and generation

Details (Click to expand)

Do not forget to activate the unifi3d virtual environment before running any script:

conda activate unifi3d # if using miniforge
# OR
# source unifi3d_venv/bin/activate # if using pyenv/venv

4.1: Benchmark a trained EncoderDecoder

Example usage (start from repository root):

python scripts/benchmarking.py benchmark=shape2vecset ae_id=ae  # ae_id is one of ae, vae or vqvae
python scripts/benchmarking.py benchmark=sdf ae_id=ae
python scripts/benchmarking.py benchmark=voxel ae_id=vae

This command takes the preprocessed ShapeNet objects, encodes and decodes them with the model, and saves the result in the outputs folder. The results include

• results.csv: Table with metric results and runtimes per sample
• mesh_comp.obj: Obj file with two rows of objects, one row showing the ground truth meshes and the back row showing the reconstructed meshes (only the first few to avoid large files)

There is a checkpoint path specified in these benchmark configs, but you can overwrite it to test your own checkpoint, e.g. with

python scripts/benchmarking.py benchmark=sdf ckpt_path=/path/to/checkpoint
python scripts/benchmarking.py benchmark=voxel ckpt_path=/path/to/checkpoint

If you used different arguments for training the model, you need to specify those as well:

python scripts/benchmarking.py benchmark=shape2vecset ckpt_path=/path/to/checkpoint net_encode.layer_norm_encoding=true

Quick testing mode

Some metrics take quite long to compute, and sometimes this script can be helpful to quickly check whether a trained model works. For testing, set the following parameters in the benchmark config:

- metrics: testing # will only compute two metrics (and sample less points to make them faster)
plot_mesh: True
limit: 10 # will only test on 10 meshes instead of the whole dataset

Collect all results and save latex tables

To collect all results from all representations, after running the benchmarking script for each of them, there is a script that load the result files and combines them:

python scripts/evaluate/make_reconstruction_table.py

This will directly print the latex code for the table. The output folders for each representation, and their name to display in the table, are currently hardcoded in the script.

4.2: Benchmark a diffusion model

Example usage (start from repo root):

python scripts/diffusion/benchmark_gen.py sample=sdf ae_id=ae diff_id=dit cat_id=Chair
python scripts/diffusion/benchmark_gen.py sample=voxel ae_id=ae diff_id=dit cat_id=Chair
python scripts/diffusion/benchmark_gen.py sample=shape2vecset ae_id=ae diff_id=dit cat_id=Chair

These sample configs include paths to the currently best checkpoints. You can use your own AE and diffusion checkpoints by setting the corresponding arguments:

python scripts/diffusion/benchmark_gen.py sample=sdf_vqvae_dit ckpt_path=/change/to/your/diffusion/model.safetensors net_encode.ckpt_path=/change/to/your/ae/model.safetensors

Outputs are saved to the output_path specified in the config file. The folder will contain the following outputs:

• results.csv: Table with runtimes and metrics per sample
• generated_samples: Folder with all the obj files for the generated meshes
• metrics.json: Json file with the distributional metrics (Cov, MMD, 1-NNA)

Computing the metrics for data that was already generated

If you already have a folder with generated samples, and just want to run the generated metrics on it, there is a script that uses a config file and only computes the distribution metrics:

python scripts/evaluate/compute_uncond_gen_metrics.py sample=sdf_vqvae_dit

Again, the metrics are saved in output_path/metrics.json

Collect all results and save latex tables

To collect all results from all representations, after running the benchmarking script for each of them, there is a script that load the result files and combines them:

python scripts/evaluate/make_generation_table.py

This will directly print the latex code for the table. The output folders for each representation, and their name to display in the table, are currently hardcoded in the script.

4.3 Special plots for the ground truth data

We show in our paper that the metrics for unconditional generation require sufficient data to compute. This experiment computes these metrics on the test set of ShapeNet and compare to the train set as a reference, which should result in perfect scores. There are two scripts for computing the unconditional metrics on the train and test split of Shapenet.

scripts/evaluate/ground_truth_uncond_metrics_compute.py
scripts/evaluate/ground_truth_uncond_metrics_plot.py

You can get a description of the parameters by calling the scripts with --help.

5: Hyperparameter optimisation

Details (Click to expand)

We use Hydra and Optuna to provide key-in-hand hyperparameter optimization capability.

• As for the training and evaluation pipelines, make sure that the unifi3d virtual environment is activated:

  conda activate unifi3d

• Setup an experiment parameter .yaml file in configs/experiment, for instance acc_voxel_vqvae.yaml. Within this file, set the parameters of your experiment.

• Setup an optuna sweeper parameter .yaml file in configs/hparams_search, for instance acc_voxel_vqvae_optuna.yaml. Within this file, set the hyperparameters you wish to optimize, for instance the learning rate or the batch size. You may use the provided template as a starting point.

• For multi-GPU hyperparameter optimisation, we rely on the hydra ray launcher pluggin to distribute the jobs -- with the parameters sampled by optuna -- to each GPU. In practice you may simply set the value of the num_parallel_process parameter to the number of available GPUs. You may use the provided parallel template as a starting point.

• Run the optuna hyperparameter optimizer:

  python scripts/accelerate/acc_train.py hparams_search=acc_voxel_vqvae_optuna_parallel experiment=acc_voxel_vqvae

  The training script should start automatically, instanciate a Trainer following the specifications of the Optuna parameter search algorithm, set in acc_voxel_vqvae_optuna_parallel.yaml.

• In case you are still in debug phase and with to see the crash logs, you may use the HYDRA_FULL_ERROR=1 before your command:

  HYDRA_FULL_ERROR=1 python scripts/accelerate/acc_train.py hparams_search=acc_voxel_vqvae_optuna_parallel experiment=acc_voxel_vqvae

6: User study

Details (Click to expand)

There are three steps:

1. Prepare the data

This requires blender 4.1 to be in the PATH

python scripts/user_study/prepare_user_study_data.py output_dir=/path/to/output_dir

2. Collect data from users

Run

python scripts/user_study/user_study.py /path/to/output_dir

Share the link

3. Compute scores

Open URL this will generate user_results/csv_data/user_study_pairwise_prefs.csv. Using the download button is not necessary.

Run

python scripts/evaluate/compute_elo_scores.py --data_dir user_results/csv_data/user_study_pairwise_prefs.csv --exp_dir XXXX --methods doctree_vae_unet_champion sdf_ae_dit voxel_ae_unet shape2vecset_ae_dit shap_e_ae_dit triplane_ae_unet

7: Contributing

Details (Click to expand)

Please check the contribution guidelines.

8: Relevant Research used in this work

Details (Click to expand)

As most research projects, Unifi3D stands on the shoulders of giants. Here are some of them:

• DualOctreeGNN
• O-CNN
• SDFusion
• Shap-E
• Convolutional Occupancy Networks

Relevant Code Frameworks

Details (Click to expand)

Part of the code was inspired by the following repositories. We thank the authors of these projects for their great work.

• Lightning-Hydra-Template
• Accelerate
• Diffusers
• Hydra
• Optuna
• pyenv
• Pytorch

9: Citation

@article{unifi3d,
  title={{Unifi3D: A Study on 3D Representations for Generation and Reconstruction in a Common Framework}},
  author={Wiedemann, Nina and Liu, Sainan and Leboutet, Quentin and Gao, Katelyn and Ummenhofer, Benjamin and Paulitsch, Michael and Yuan, Kai},
  journal={Transactions on Machine Learning Research},
  year={2025},
  url={https://openreview.net/forum?id=GQpTWpXILA},
}