Mini-Batch Consistent Slot Set Encoder forScalable Set Encoding

This repository contains the accompanying code for SSE.

Dependencies

torch==1.8.0
torchvision==0.9.0
torch-cluster==1.5.9
torch-geometric==1.7.0
torch-scatter==2.0.6
torch-sparse==0.6.9
torch-spline-conv==1.2.1

Empirical Verification of Proposition 1 & 2

To empirically verify Propositions 1 & 2, run the following:

python models/slotsetencoder.py

This will check three things: permutation equivariance, permutation invariance and Mini-Batch Consistency. The test code verifies these for the Slot Set Encoder with both random and deterministic slot initialization.

Training

Run the following files to reproduce the results on CelebA image reconstruction.

bash scripts/celeba/deepsets/training/minibatch_training/_setsize_sh

bash scripts/celeba/settransformer/training/minibatch_training/_setsize_.sh

bash scripts/celeba/consistent/training/minibatch_training_xxx/_setsize_.sh

_setsize_ is one of 100, 200, 500 and 1000. For the Slot Set Encoder model, xxx is either learned or random.

Mini-Batch Testing for MBC models.

After training, you can perform Mini-Batch testing with the model trained on 1000 elements by running:

bash scripts/celeba/deepsets/testing/generalization/1000.sh
bash scripts/celeba/deepsets/testing/generalization/1000_small.sh

bash scripts/celeba/consistent/testing/generalization_xxx/1000.sh
bash scripts/celeba/consistent/testing/generalization_xxx/1000_small.sh