Code to accompany the paper A General Method for Amortizing Variational Filtering by Marino et al., NeurIPS, 2018.
First, clone the repository by opening a terminal and running:
$ git clone https://github.com/joelouismarino/amortized-variational-filtering
Then enter the project directory:
$ cd amortized-variational-filtering
To avoid dependency conflicts, create a conda environment using Anaconda, which you can download here. Once you have installed Anaconda, use the environment file, avf_env.yml
, to create the environment and install the dependencies:
$ conda env create -f avf_env.yml
To activate the environment, run:
$ conda activate avf
The terminal should now appear as follows:
(avf) $
Within the environment, install PyTorch by visiting the list of versions here, and installing version 0.3.0.post4
for your version of CUDA (8.0
or 9.0
). For example, with CUDA 9.0
, you would run:
(avf) $ pip install https://download.pytorch.org/whl/cu90/torch-0.3.0.post4-cp27-cp27mu-linux_x86_64.whl
Note: if you want to use the BAIR Robot Pushing dataset, you will also need to install tensorflow. At this point, you should be able to run code within the AVF repository. After running the code, when you are ready to exit the environment, run:
(avf) $ conda deactivate
or simply close out of the terminal window.
The code uses visdom to plot various metrics during training. To start visdom, open a terminal window, activate the environment (see above), and run:
(avf) $ python -m visdom.server
You can now open a web browser and navigate to https://localhost:8097/
to view the visdom output. Each experiment is saved as separate visdom environment using the time stamp of when the training run was started.
Each experiment is defined by a set of configuration dictionaries, which are found in the config
directory. These consist of train_config
, run_config
, data_config
, and model_config
. The train_config
dictionary contains training specifications, such as the batch size, learning rate, etc. The run_config
dictionary contains specifications for the current run, such as which GPU to use, where to log the results, etc. The data_config
dictionary contains specifications for the data, like the dataset and sequence length. Finally, the model_config
contains the specifications for the generative and inference models.
You can define an experiment by altering each of the config files in the config
directory. Alternatively, you can use one of the pre-defined configurations from the paper, which are located in the experiments
directory. We will now walk through how to run the code with each of these approaches.
To run the code with your own set of configuration parameters (located in the config
directory), run:
(avf) $ python train.py
To instead run one of the pre-defined experiment configurations, use the command line arguments --dataset
, --model
, and --inference
. For example, to run SVG on the KTH Actions dataset with AVF, run:
(avf) $ python train.py --dataset 'kth_actions' --model 'svg' --inference 'avf'
To run the baseline filtering method (proposed in Denton et al., 2018), run:
(avf) $ python train.py --dataset 'kth_actions' --model 'svg' --inference 'baseline'
This will load configuration parameters for model_config
and train_config
from the corresponding experiment in the experiments
directory. Note that you will still need to define config/run_config.py
, as well as data_path
in config/data_config.py
. If the dataset has not been downloaded yet, the code will automatically download the data to the data_path
. Note that not all models can be run with all datasets. The TIMIT dataset must be downloaded manually here.
Evaluation for some datasets is performed on each test sequence separately, i.e. not cropping out sub-sequences. Since test sequences can be of various lengths, you have to perform evaluation with a batch_size
(see config/train_config.py
) of 1. To evaluate a model, specify the resume_path
in config/run_config.py
, then run:
(avf) $ python evaluate.py
This will run evaluation and save the results using pickle
.
The repository is set-up to be fairly straightforward to extend to new datasets and new models. We'll now briefly describe how to do each of these.
The datasets are defined in util/data/datasets/
. To add a new dataset, create a class for the dataset in this directory and import the dataset in util/data/datasets/__init__.py
. For example, you could define a dataset called MyDataset
in a file util/data/datasets/my_dataset.py
:
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, args):
# define initialization here
def __getitem__(self, index):
# define getitem method here
def __len__(self):
# define dataset length method here
Then, in util/data/datasets/__init__.py
, include the following line:
from my_dataset import MyDataset
The datasets are loaded in util/data/load_dataset.py
. Add another elif
statement in this file for your dataset and create train
, val
, and test
datasets. For example,
elif dataset_name == 'my_dataset':
# do whatever downloading, etc. you need to do here
from datasets import MyDataset
train = MyDataset(train_args)
val = MyDataset(val_args)
test = MyDataset(test_args)
Note that you may want to include data transforms in the dataset (e.g. random sequence crops, image crops/flips, normalization, etc.). See the other datasets in load_dataset.py
for examples of how to include these. Once you have added your dataset to load_dataset.py
, you can use your dataset by specifying the name for dataset_name
in config/data_config.py
. You can also specify the data_type
to add additional arguments to data_config
. When in doubt, use other
. You will also need to specify the length of sequences for training. Here, we set this to 10
. For example,
data_config = {
'data_path': '/path/to/datasets/',
'dataset_name': 'my_dataset',
'data_type': 'other', # video, audio, other, etc.
'sequence_length': 10
}
Note that the current set of models have different versions for various datasets, e.g. based on the number of data dimensions, etc. You can specify which version of the model to use in config/model_config.py
, and/or define a new version of the model in the corresponding model file found in lib/models/
.
To add a new model to the repository, you will need to define it in lib/models/
as a separate file. All models inherit from the base class LatentVariableModel
defined in lib/models/latent_variable_model.py
, which defines the set of methods that are expected in the code. These required methods are
infer
: perform inference for a given data input stepgenerate
: generate a prediction/reconstruction from the modelstep
: step the latent dynamics forward one stepre_init
: re-initialize the model's stateinference_parameters
: returns the parameters of the inference modelgenerative_parameters
: returns the parameters of the generative modelinference_mode
: puts the model into inference mode, only backpropagating into the inference modelgenerative_mode
: puts the model into generative mode, backpropagating through the generative model.
You can also internally define whatever other methods you need for your model. For instance, the models are currently implemented with an internal _construct
method. To create a model called MyModel
, you would create a new file in lib/models
called my_model.py
with something like the following:
from latent_variable_model import LatentVariableModel
class MyModel(LatentVariableModel):
def __init__(self, model_config):
super(MyModel).__init__(model_config)
# initialize the model here
def infer(self, observation):
# define inference method here
... # define the rest of the required methods (see above)
Then, in the file lib/models/load_model.py
, add another elif
statement for your model:
elif model_name == 'my_model':
from my_model import MyModel
return MyModel(model_config)
Finally, to use your model, specify your model in the architecture
field in config/model_config.py
. For instance:
model_config = {
'architecture': 'my_model',
...
}
If you have any questions, feel free to send me an email at jmarino
[at] caltech
[dot] edu
, or post an issue on Github.