This repository contains (a subset of) our code from our preprint: Disentangling Fact from Grid Cell Fiction in Trained Deep Path Integrators.
As background, our NeurIPS 2022 paper No free lunch from deep learning in neuroscience: A case study through models of the entorhinal-hippocampal circuit was unable to reproduce key findings of Sorscher et al. NeurIPS 2019 A unified theory for the origin of grid cells through the lens of pattern formation. After the NeurIPS 2022 review process, code subsequently released by Sorscher et al.; however, it was implemented in a deprecated version of TensorFlow (1.14.0), so we reproduced it here using PyTorch and PyTorch Lightning, then integrated W&Bs to sweep hyperparameters and log results for further analysis.
Create a conda environment with the required packages:
conda env create --file environment.yml
To activate the environment:
conda activate sorscher_reproduction
Upgrade pip:
pip install --upgrade pip
Then install wandb with pip, recommended by W&B here.
pip install wandb
Then install pytorch lightning:
pip install lightning
Then install the correct version of Pytorch:
pip install torch==1.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
I also needed to install an additional system component via conda:
conda install libgcc
pip3 install torch torchvision torchaudio
To run a single experiment, run the following command:
python train_one.py
The hyperparameters for this run are specified in src/globals.py
. The results should
automatically be logged to W&B.
To create a sweep, run the following command at the CLI:
wandb sweep <path to sweep e.g. sweeps/20230523/basic.yaml>
This will output the newly created sweep ID, e.g.,:
wandb: Creating sweep from: sweeps/20230522/basic.yaml
wandb: Creating sweep with ID: ox6goel1
wandb: View sweep at: https://wandb.ai/rylan/sorscher-2022-reproduction/sweeps/ox6goel1
wandb: Run sweep agent with: wandb agent rylan/sorscher-2022-reproduction/ox6goel1
Then create an "Agent" which will pull the next run off the queue and run it to completion, before repeating. An agent will continue to function until no more runs in the sweep remain. To create an agent, run:
wandb agent <your W&B username e.g. rylan>/sorscher-2022-reproduction/<sweep ID e.g. ox6goel1>
A sensible way to parallelize is to launch multiple jobs on a cluster, where each job contains 1 agent. So for instance, if we're running on a SLURM cluster, each SLURM job should create 1 agent. Rylan has code for this.
For analysis code, see notebooks/
.
Questions? Comments? Interested in collaborating? Open an issue or email Rylan Schaeffer ([email protected]) and cc Mikail Khona ([email protected]) and Ila Fiete ([email protected]).