Code for NeurIPS 2019 paper titled NAOMI: Non-Autoregressive Multiresolution Sequence Imputation

Code is written with PyTorch v0.4.1 (Python 3.6.5). Billiards data can be downloaded here, basketball data is available from STATS.

First open visdom, then adjust hyperparameters in train_model.sh and run the shell file.

• --model: “NAOMI” or “SingleRes”

• --task: “basketball” or “billiard”

• --y_dim: 10 for basketball and 2 for billiard

• --rnn_dim and --n_layers: gru cell size for all models, including forward and backward rnns

• --dec1_dim to --dec16_dim: For NAOMI, these values correspond to dimensions of different decoders. For SingleRes, only dec1_dim is used for decoder.

• --pre_start_lr: initial learning rate for supervised pretrain

• --pretrain: supervised pretrain epochs

• --highest: largest stepsize for NAOMI decoders, should be 2^n

• --discrim_rnn_dim and --discrim_layers: discriminator rnn size

• --policy_learning_rate: learning rate for generator in adversarial training

• --discrim_learning_rate: learning rate for discriminator in adversarial training

• --pretrain_disc_iter: number of iterations to pretrain discriminator

• --max_iter_num: number of adversarial training iterations

If you find this repository, e.g., the code and the datasets, useful in your research, please cite the following paper:

@inproceedings{liu2019naomi,
  title={NAOMI: Non-Autoregressive Multiresolution Sequence Imputation},
  author={Liu, Yukai and Yu, Rose and Zheng, Stephan and Zhan, Eric and Yue, Yisong},
  booktitle={Advances in Neural Information Processing Systems(NeurIPS '19)},
  year={2019}
}