Skip to content

Latest commit

 

History

History
102 lines (76 loc) · 4.48 KB

README.md

File metadata and controls

102 lines (76 loc) · 4.48 KB
Raw

ADEPT-Model-Release

PyTorch implementation of Approximate Derenderer, Extended Physics, and Tracking (ADEPT).

Modeling Expectation Violation in Intuitive Physics with Coarse Probabilistic Object Representations

Kevin Smith*, Lingjie Mei*, Shunyu Yao, Jiajun Wu, Elizabeth S. Spelke, Joshua B. Tenenbaum, Tomer Ullman (* indicates equal contribution)

NeurIPS 2019, CogSci 2020

Paper BibTeX Website

For data preparation, see our dataset generation repository for dataset generation.

Getting started

Prerequisites

  • Linux
  • Python3
  • PyTorch 1.1 with CUDA support
  • Other required python packages specified by requirements.txt.

Installation

  1. Clone this repository

    git clone https://github.com/JerryLingjieMei/ADEPT-Model-Release
cd ADEPT-Model-Release

    and change CONTENT_FOLDER in utils.constants to the folder of this repository.

  2. Create a conda environment for ADEPT Model and install the requirements.

    conda create --n adept-model
conda activate adept-model
pip install -r requirements.txt
conda install pytorch=1.1.0 cuda100 -c pytorch #Assume you use cuda version 10.0

  3. Set up Mask R-CNN implementation by Facebook Research, see here

  4. Change TRAIN_ROOT and HUMAN_ROOT in tools/paths_catalog.py to reflect the directory storing training videos and human test videos, and ANNOTATION_FOLDER in experiments/default_physics.yaml to the directory storing human test videos.

  5. Download pretrained weights for detection and derender by running

    . scripts/download_pretrained_model.sh

  6. (Optional) If you have multiple machines, you may change get_host_id in utils/misc.py to reflect the id of your machine. With that in hand, you may speed up physical reasoning process by using --stride N arguments, where you have N machines with consecutive ids.

Experiments

  1. Collect COCO annotations for the training set, and the human test set:

    # Fill in the TRAIN_SET_FOLDER with the training set folder
python3 dataset/makers/make_coco.py -i ${TRAIN_SET_FOLDER} -o data/ann.json
# Fill in the HUMAN_TEST_SET_FOLDER with the human set folder
python3 dataset/makers/make_coco.py -i ${HUMAN_TEST_SET_FOLDER} -o data/human_ann.json

    Then train the detection part of perception on the training set and test on training set and human test set,

    # Training and testing with multiple GPUs
export CUDA_VISIBLE_DEVICES=0,1,2,3
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/detection_train_net.py --config_file experiments/default_detection.yaml
python -m torch.distributed.launch --nproc_per_node=$NGPUS tools/detection_test_net.py --config_file experiments/default_detection.yaml
# Training and testing on a single GPU
python tools/detection_train_net.py --config_file experiments/default_detection.yaml
python tools/detection_test_net.py --config_file experiments/default_detection.yaml

  2. Collect Object Proposal annotations for the training set and train the approximate derenderer part of the perception.

    # Fill in the TRAIN_SET_FOLDER with the training set folder
python3 dataset/makers/make_proposal.py -i ${TRAIN_SET_FOLDER} -o data/annotated_ann.json -s output/default_detection/inference/physics/segm.json
python3 tools/train_net.py --config_file experiments/default_derender.yaml

    Collect Object Proposal annotations for the human test set and test the approximate derender.

    # Fill in the HUMAN_TEST_SET_FOLDER with the human set folder
python3 dataset/makers/make_proposal.py -i ${HUMAN_TEST_SET_FOLDER} -o data/annotated_human_ann.json -s output/default_detection/inference/physics_human/segm.json
python3 tools/test_net.py --config_file experiments/default_derender.yaml

  3. Evaluate the physical reasoning result on the inferred objects proposals of human test set:

    # Single machine
python3 tools/run_physics.py --config_file experiments/default_physics.yaml
# Multiple (e.g. 8) machines
python3 tools/run_physics.py --config_file experiments/default_physics.yaml --stride 8 #On each machine