Code for "V1T: Large-scale mouse V1 response prediction using a Vision Transformer"
MIT License
Code for TMLR2023 paper "V1T: Large-scale mouse V1 response prediction using a Vision Transformer".
Authors: Bryan M. Li, Isabel M. Cornacchia, Nathalie L. Rochefort, Arno Onken
@article{
li2023vt,
title={V1T: large-scale mouse V1 response prediction using a Vision Transformer},
author={Bryan M. Li and Isabel Maria Cornacchia and Nathalie Rochefort and Arno Onken},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2023},
url={https://openreview.net/forum?id=qHZs2p4ZD4},
note={}
}
We sincerely thank Willeke et al. for organizing the Sensorium challenge and, along with Franke et al., for making their high-quality large-scale mouse V1 recordings publicly available. This codebase is inspired by sinzlab/sensorium, sinzlab/neuralpredictors and sinzlab/nnfabrik.
The codebase repository has the following structure. Check .gitignore for the ignored files.
sensorium2022/
data/
sensorium/
static21067-10-18-GrayImageNet-94c6ff995dac583098847cfecd43e7b6.zip
...
franke2022/
static25311-4-6-ColorImageNet-104e446ed0128d89c639eef0abe4655b.zip
...
README.md
misc/
src/
v1t/
...
.gitignore
README.md
setup.sh
demo.ipynb
submission.py
sweep.py
train.py
...
demo.ipynb
demonstrates how to load the best V1T model and inference the Sensorium+ test set, as well as extracting the attention rollout maps.sweep.py
performs hyperparameter tuning using Weights & Biases.train.py
contains the model training procedure.conda create -n v1t python=3.10
v1t
virtual environment
conda activate v1t
setup.sh
script to install the relevant conda
and pip
packages for macOS and Ubuntu devices.
sh setup.sh
# install PyTorch
conda install -c pytorch pytorch=2.0 torchvision torchaudio -y
# install V1T package
pip install -e .
python train.py --dataset data/sensorium --output_dir runs/v1t_model --core vit --readout gaussian2d --behavior_mode 3 --batch_size 16
--help
flag to see all available optionstrain.py
uses both TensorBoard and Weights & Biases to log training information.
tensorboard --logdir runs/v1t_model --port 6006
localhost:6006
on your browser--use_wandb
and (optional) --wandb_group <group name>
to enable wandb
logging.