This is an official PyTorch implementation of the SPACE model presented in the following paper:
SPACE: Unsupervised Object-Oriented Scene Representation via Spatial Attention and Decomposition
{Zhixuan Lin, Yi-Fu Wu, Skand Vishwanath Peri}, Weihao Sun, Gautam Singh, Fei Deng, Jindong Jiang, Sungjin Ahn
ICLR 2020
Project page
Project directories:
src: source codedata: where you should put the datasetsoutput: anything the program outputs will be saved here. These includeoutput/checkpoints: training checkpoints. Also, model weights with the best performance will be saved hereoutput/logs: tensorboard event filesoutput/eval: quantitative evaluation resultsoutput/demo: demo images
scripts: some useful scripts for downloading things and showing demospretrained: where to put downloaded pretrained models
This project uses YACS for managing experiment configurations. Configurations are specified with YAML files. These files are in src/configs. We provide five YAML files that correspond to the figures in the paper:
3d_room_large.yaml: for the 3D Room Large dataset3d_room_small.yaml: for 3D Room Small datasetatari_spaceinvaders.yaml: for the Space Invaders gameatari_riverraid.yaml: for the River Raid gameatari_joint.yaml: for joint training on 10 Atari games
This project uses Python 3.7 and PyTorch 1.3.0.
Create a conda environment with Python 3.7 and activate it. Other versions of Python should also be fine:
conda create -n space python=3.7
conda activate space
Install PyTorch 1.3.0:
pip install torch==1.3.0+cu100 torchvision==0.4.1+cu100 -f https://download.pytorch.org/whl/torch_stable.html
Note that this requires CUDA 10.0. If you need CUDA 9.2 then change cu100 to cu92. Depending on your cuda version, you may want to install previous versions of PyTorch. See here.
Other requirements are in requirements.txt and can be installed with
pip install -r requirements.txt
TensorBoard is used for training visualization and included in requirements.txt. In some cases, TensorBoard without a full TensorFlow installation can show some weird behaviors (not loading logs, etc.). If that happens, consider installing full TensorFlow.
The following datasets with Google Drive download links are provided:
- 3D Room Large :
OBJ3D_LARGE.tar.gz(1.4G) - 3D Room Small:
OBJ3D_SMALL.tar.gz(156M) - Atari (with 11 games):
ATARI.tar.gz(2.2G)
Depending on your need, you can download one or more of these datasets. Two download options are available:
-
Download with scripts. Run one or more of the following scripts:
# Run one or more of these sh scripts/download_data_3dlarge.sh sh scripts/download_data_3dsmall.sh sh scripts/download_data_atari.shDatasets will be downloaded to
dataand decompressed. -
Manual download. Alternatively, you can manually download them with the provided Google Drive links, put them under the
datadirectory, and decompress them withtar -xzf [FILE].tar.gz.
The data directory should look like this (if you have downloaded all three datasets):
data
├── OBJ3D_LARGE
│ ├── test
│ ├── train
│ └── val
├── OBJ3D_SMALL
│ ├── test
│ ├── train
│ └── val
└── ATARI
├── Asterix-v0
├── Atlantis-v0
├── Carnival-v0
├── DoubleDunk-v0
├── Kangaroo-v0
├── MontezumaRevenge-v0
├── MsPacman-v0
├── Pooyan-v0
├── Qbert-v0
├── Riverraid-v0
└── SpaceInvaders-v0
To download pretrained models, two options are available:
-
Download with scripts. Run the following script to download pretrained models:
sh scripts/download_pretrained.shPretrained models will be downloaded to the
pretraineddirectory and decompressed. -
Manual download. Alternatively, you can manually download the files (in one compressed file
pretrained.tar.gz) with this Google Drive link (239M), put it under thepretraineddirectory and decompress it withtar -xzf pretrained.tar.gz.
The pretrained directory should look like this:
pretrained
├── 3d_room_large.pth
├── 3d_room_small.pth
├── atari_spaceinvaders.pth
├── atari_riverraid.pth
└── atari_joint.pth
Then run one of the following to create some visualizations, with either CPU or GPU:
# Use CPU
sh scripts/demo_cpu.sh
# Use GPU
sh scripts/demo_gpu.sh
Images showing foreground objects and background segmentation will be dumped to output/demo. If you have downloaded all three datasets, then five images 3d_room_large.png, 3d_room_small.png, atari_spaceinvaders.png, atari_riverraid.png and atari_joint.png will be generated in output/demo. Otherwise only some of them will be generated.
If you are using a remote server, you can then run
python -m http.server -d output/demo 8080
and then visit http://[your server's address]:8080 in your local browser to view these images.
First, cd src. Make sure you are in the src directory for all commands in this section. All paths referred to are also relative to src.
The general command to run the program is (assuming you are in the src directory)
python main.py --task [TASK] --config [PATH TO CONFIG FILE] [OTHER OPTIONS TO OVERWRITE DEFAULT YACS CONFIG...]
Detailed instructions will be given below.
Training. Run one or more of the following to train the model on the datasets you want:
-
3D Room Large:
python main.py --task train --config configs/3d_room_large.yaml resume True device 'cuda:0' -
3D Room Small:
python main.py --task train --config configs/3d_room_small.yaml resume True device 'cuda:0' -
River Raid:
python main.py --task train --config configs/atari_riverraid.yaml resume True device 'cuda:0' -
Space Invaders:
python main.py --task train --config configs/atari_spaceinvaders.yaml resume True device 'cuda:0' -
Joint training on 10 Atari games:
python main.py --task train --config configs/atari_joint.yaml resume True device 'cuda:0'
These start training with GPU 0 (cuda:0). There some useful options that you can specify. For example, if you want to use GPU 5, 6, 7, and 8 and resume from checkpoint ../output/checkpoints/3d_room_large/model_000008001.pth, you can run the following:
python main.py --task train --config configs/3d_room_large.yaml \
resume True resume_ckpt '../output/checkpoints/3d_room_large/model_000008001.pth' \
parallel True device 'cuda:5' device_ids '[5, 6, 7, 8]'
Other available options are specified in config.py.
Training visualization. Run the following
# Run this from the 'src' directory
tensorboard --bind_all --logdir '../output/logs' --port 8848
And visit http://[your server's address]:8848 in your local browser.
Evaluation. We only have ground truth for the two 3D Room datasets. After training is finished with the specified maximum steps (or you can stop them manually), run the following to evaluate APs and counting accuracy:
# Run one or more of these
python main.py --task eval --config configs/3d_room_large.yaml resume True device 'cuda:0'
python main.py --task eval --config configs/3d_room_small.yaml resume True device 'cuda:0'
The model with the best performance (average AP) will be loaded and evaluated. The results will be printed to stdout. They will also be saved to ../output/eval as JSON files.
Alternatively, you can directly run evaluation using the pretrained models that we provided:
# Run one or more of these
python main.py --task eval --config configs/3d_room_large.yaml resume True device 'cuda:0' resume_ckpt '../pretrained/3d_room_large.pth'
python main.py --task eval --config configs/3d_room_small.yaml resume True device 'cuda:0' resume_ckpt '../pretrained/3d_room_small.pth'
- For some reason we were using BGR images for our Atari dataset and our pretrained models can only handle that. Please convert the images to BGR if you are to test your own Atari images with the provided pretrained models.
- There is a chance that SPACE doesn't learn proper background segmentation for the 3D Room Large datasets. Due to the known PyTorch reproducibity issue, we cannot guarantee each training run will produce exactly the same result even with the same seed. For the 3D Room Large datasets, if the model doesn't seem to be segmenting the background in 10k-15k steps, you may considering changing the seed and rerun (or not even changing the seed, it will be different anyway). Typically after trying 1 or 2 runs you will get a working version.
If you want to apply SPACE to your own task (e.g., for RL), please be careful. Applying SPACE to RL is also our original intent, but we found that the model can sometimes be unstable and sensitive to hyperparameters and training tricks. There are several reasons:
- The definition of objects and background is ambiguous in many cases. Atari is one case where objects are often well-defined. But in many other cases, it is not. For more complicated datasets, making SPACE separate foreground and background properly can be something non-trivial.
- Learning is difficult when object sizes vary a lot. In SPACE, we need to set a proper prior for object sizes manually and that turn out to be crucial hyperparameter. For example, for the 10 Atari games we tested, objects are small and roughly of the same size. When object sizes vary a lot SPACE may fail.
That said, we are pleased to offer discussions and pointers if you need help (especially when fine-tuning it on your own dataset). We also hope this will facilitate future works that overcome these limitations.
If you find this code useful for your research, please cite our paper with the following BibTeX entry
@inproceedings{
Lin2020SPACE,
title={SPACE: Unsupervised Object-Oriented Scene Representation via Spatial Attention and Decomposition},
author={Zhixuan Lin and Yi-Fu Wu and Skand Vishwanath Peri and Weihao Sun and Gautam Singh and Fei Deng and Jindong Jiang and Sungjin Ahn},
booktitle={International Conference on Learning Representations},
year={2020},
url={https://openreview.net/forum?id=rkl03ySYDH}
}
The evaluation code is adapted from the one used in SPAIR. The code structure is inspired (and significantly simplified) by Mask-RCNN (deprecated, with the latest being Detectron2) from Facebook. Google Drive download commands are created with https://gdrive-wget.glitch.me/