Code to reproduce the trained RNN in a unified theory for the origin of grid cells through the lens of pattern formation (NeurIPS '19) and additional analysis described in this preprint.
Quick start:
-
inspect_model.ipynb: Train a model and visualize its hidden unit ratemaps.
-
main.py: or, train a model from the command line.
-
pattern_formation.ipynb: Numerical simulations of pattern-forming dynamics.
Includes:
-
trajectory_generator.py: Generate simulated rat trajectories in a rectangular environment.
-
place_cells.py: Tile a set of simulated place cells across the training environment.
-
model.py: Contains the vanilla RNN model architecture, as well as an LSTM.
-
trainer.py: Contains model training loop.
-
models/example_trained_weights.npy Contains a set of pre-trained weights.
We recommend creating a virtual environment:
$ virtualenv env
$ source env/bin/activate
$ pip install --upgrade pipThen, install the dependencies automatically with pip install -r requirements.txt
or manually with:
$ pip install --upgrade numpy==1.17.2
$ pip install --upgrade tensorflow==2.0.0rc2
$ pip install --upgrade scipy==1.4.1
$ pip install --upgrade matplotlib==3.0.3
$ pip install --upgrade imageio==2.5.0
$ pip install --upgrade opencv-python==4.1.1.26
$ pip install --upgrade tqdm==4.36.0
$ pip install --upgrade opencv-python==4.1.1.26
$ pip install --upgrade torch==1.10.0If you want to train your own models, make sure to properly set the default
save directory in main.py!