Compositional Scene Understanding through Inverse Generative Modeling
Yanbo Wang 1,
Justin Dauwels 1,
Yilun Du 2
1TU Delft, 2Harvard
We propose a compositional inverse generative modeling (IGM) framework for visual concept inference with zero-shot generalization capabilities. Our approach infers visual concepts by optimizing the conditional parameters of a compositional generative model, enabling interpreting images that are more complex than those seen during training.
Run the following to create and activate a conda environment:
conda env create -f environment.yml
conda activate IGM
To train a compositional generative model on CLEVR, run the following:
CUDA_VISIBLE_DEVICES=0,1 accelerate launch --multi_gpu --num_processes=2 \
--main_process_port 28100 train_clevr.py --enable_xformers_memory_efficient_attention \
--dataloader_num_workers 4 --learning_rate 2e-5 --mixed_precision fp16 --num_validation_images 10 \
--val_batch_size 10 --max_train_steps 500000 --checkpointing_steps 10000 --checkpoints_total_limit 1 \
--gradient_accumulation_steps 1 --seed 42 \
--output_dir ./outputs/ \
--scheduler_config configs/clevr-2D-pos/scheduler/scheduler_config.json \
--unet_config configs/clevr-2D-pos/unet/config.json \
--dataset_root /space/ywang86/CLEVR_v1.0 --train_batch_size 128 \
--resolution 64 --validation_steps 1000 --tracker_project_name clevr
To train a compositional generative model on CelebA, run the following:
CUDA_VISIBLE_DEVICES=0,1 accelerate launch --multi_gpu --num_processes=2 \
--main_process_port 28100 train_celeba.py --enable_xformers_memory_efficient_attention \
--dataloader_num_workers 4 --learning_rate 2e-5 --mixed_precision fp16 --num_validation_images 10 \
--val_batch_size 10 --max_train_steps 50000 --checkpointing_steps 10000 --checkpoints_total_limit 10 \
--gradient_accumulation_steps 1 --seed 42 \
--output_dir ./outputs \
--scheduler_config configs/celeba/scheduler/scheduler_config.json \
--unet_config configs/celeba/unet/config.json \
--dataset_root /space/ywang86/celeba_labels --train_batch_size 64 \
--resolution 128 --validation_steps 1000 --tracker_project_name celeba
To infer object locations, run the following:
CUDA_VISIBLE_DEVICES=0 python infer_clevr.py \
--learning_rate 4e-3 \
--scheduler_config configs/clevr-2D-pos/scheduler/scheduler_config.json \
--unet_config configs/clevr-2D-pos/unet/config.json \
--dataset_root /space/ywang86/CLEVR_v1.0 --resolution 64
To infer facial attributes, run the following:
CUDA_VISIBLE_DEVICES=0 python infer_celeba.py \
--scheduler_config configs/clevr-2D-pos/scheduler/scheduler_config.json \
--unet_config configs/clevr-2D-pos/unet/config.json \
--dataset_root /space/ywang86/celeba_labels
To infer animal categories with pretrained SD, run the following:
CUDA_VISIBLE_DEVICES=0 python infer_animal.py
See our paper for details on training datasets.
| Dataset | Link |
|---|---|
| CLEVR | Link |
| CelebA | Link |
| Animal | Provided in the data folder |
