This is the pytorch code for the paper Unsupervised Learning of Compositional Energy Concepts.

Please download a pretrained model at this link and then execute the following code to test a pretrained CelebA-HQ 128x128 COMET model

python demo.py --im_path=im0.jpg

Please utilize the following command to run global factor decomposition on CelebA-HQ (or other datasets)

python train.py --exp=celebahq --batch_size=12 --gpus=1 --cuda --train --dataset=celebahq --step_lr=500.0

You may further run the code on high-resolution 128x128 images below

python train.py --exp=celebahq_128 --batch_size=12 --gpus=1 --cuda --train --dataset=celebahq_128 --step_lr=500.0

Please utilize the following command to run local factor decomposition on CLEVR

python train.py --exp=clevr_local_decomp --num_steps=5 --step_lr=1000.0 --components=5 --dataset=clevr --cuda --train --batch_size=24 --latent_dim=16 --recurrent_model --pos_embed

Please utilize the following link to download the CLEVR dataset utilized in our experiments. Downloads for additional datasets will be posted soon. Feel free to raise an issue if there is a particular dataset you would like downloaded

If you find our code useful for your research, please consider citing

@inproceedings{du2021comet,
  title={Unsupervised Learning of Compositional Energy Concepts},
  author={Du, Yilun and Li, Shuang and Sharma, Yash and Tenenbaum, B. Joshua
  and Mordatch, Igor},
  booktitle={Advances in Neural Information Processing Systems},
  year={2021}
}