Unsupervised Resource Allocation with Graph Neural Networks
Check out the paper here.
PyTorch code for the forward model of our algorithm can be found in this repository in the file
model.py. To train the model, execute
Data required to train this model can be found here
Requirements for our codebase can be found in
environment.yml. Note that one needs to use the following custom astropy:
pip install git+https://github.com/MilesCranmer/astropy
(it has some of the Cosmology calculations vectorized).
If you are using
conda, and have CUDA version 11.0 and cuDNN version 8.0, you can create a duplicate of our env, using:
which will create a new environment called
gnn_allocation. This uses PyTorch 1.7.1, though it is likely to work for other versions if you decide to modify
environment.yml. You can also use an implementation without CUDA using the