|
| 1 | +# Algonauts 2023 |
| 2 | + |
| 3 | +Code for the [CMI-DAIR submission](https://arxiv.org/abs/2308.02351) to the [Algonauts 2023 Challenge](http://algonauts.csail.mit.edu/) (team "BlobGPT"). |
| 4 | + |
| 5 | +<p align="center"> |
| 6 | + <img src=".github/arch.svg" alt="model architecture" width="360"> |
| 7 | +</p> |
| 8 | + |
| 9 | +Our model consists of a multi-subject linear encoding head attached to a pretrained trunk model. The multi-subject head has three components: (1) a shared multi-layer feature projection, (2) shared plus subject-specific low-dimension linear transformations, and (3) a shared frozen PCA embedding. The feature projection is "factorized" as a 1x1 convolution followed by learned depthwise spatial pooling. |
| 10 | + |
| 11 | +Our submission model used an `eva02_base_patch14_224.mim_in22k` trunk from [timm](https://github.com/huggingface/pytorch-image-models). We first trained the encoding head only with the trunk frozen (phase 1). Then we unfroze the trunk's attention blocks and fine-tuned the model end-to-end (phase 2). See [our report](https://arxiv.org/abs/2308.02351) for more details. |
| 12 | + |
| 13 | +## Results |
| 14 | + |
| 15 | +| Model | Val score | Test score | Config | Weights | |
| 16 | +| --- | --- | --- | --- | --- | |
| 17 | +| GroupLin-P1 | 20.4% | 58.8% | [config](config/phase1_head_only.yaml) | [weights](https://github.com/cmi-dair/algonauts23/releases/download/v0.1.0/grouplin_phase1.pt) | |
| 18 | +| GroupLin-P2 | 20.9% | 60.3% | [config](config/phase2_finetune.yaml) | [weights](https://github.com/cmi-dair/algonauts23/releases/download/v0.1.0/grouplin_phase2.pt) | |
| 19 | + |
| 20 | +**Val score**: median R<sup>2</sup> on our validation set |
| 21 | + |
| 22 | +**Test score**: mean noise-normalized R<sup>2</sup> on official challenge test set |
| 23 | + |
| 24 | +## Installation |
| 25 | + |
| 26 | +Clone the repository. |
| 27 | + |
| 28 | +```bash |
| 29 | +git clone https://github.com/cmi-dair/algonauts23.git |
| 30 | +``` |
| 31 | + |
| 32 | +Create a new environment. |
| 33 | + |
| 34 | +```bash |
| 35 | +cd algonauts23 |
| 36 | +python3.10 -m venv --prompt algonauts23 .venv |
| 37 | +source .venv/bin/activate |
| 38 | +pip install -U pip |
| 39 | +``` |
| 40 | + |
| 41 | +Install the dependencies |
| 42 | + |
| 43 | +```bash |
| 44 | +pip install -r requirements.txt |
| 45 | +``` |
| 46 | + |
| 47 | +Install the package |
| 48 | + |
| 49 | +```bash |
| 50 | +pip install . |
| 51 | +``` |
| 52 | + |
| 53 | +## Dataset preparation |
| 54 | + |
| 55 | +Follow the steps [here](dataset/) to download and prepare the data for training. |
| 56 | + |
| 57 | +## Training |
| 58 | + |
| 59 | +### WandB setup (optional) |
| 60 | + |
| 61 | +We used [WandB](https://wandb.ai/) for experiment tracking. To set up WandB, create an account if you don't have one already and get your [API key](https://docs.wandb.ai/quickstart#common-questions). Then include these commands before launching training. |
| 62 | + |
| 63 | +```bash |
| 64 | +export WANDB_API_KEY="XXXXXXX" |
| 65 | +wandb login |
| 66 | +opts="--wandb" |
| 67 | +``` |
| 68 | + |
| 69 | +### Download PCA weights |
| 70 | + |
| 71 | +Our model uses a frozen group PCA embedding. You can download the weights [here](https://github.com/cmi-dair/algonauts23/releases/download/v0.1.0/group_pca_d-2048.pt). |
| 72 | + |
| 73 | +You can also re-run the group PCA using [`scripts/fit_group_pca.py`](scripts/fit_group_pca.py). |
| 74 | + |
| 75 | +### Phase 1: Multi-subject linear head only |
| 76 | + |
| 77 | +In the first phase of training, we only train the multi-subject linear head with the trunk model frozen. |
| 78 | + |
| 79 | +```bash |
| 80 | +python scripts/train_group_encoder.py config/phase1_head_only.yaml \ |
| 81 | + --out_dir results --workers 4 $opts |
| 82 | +``` |
| 83 | + |
| 84 | +### Phase 2: Partial fine-tuning |
| 85 | + |
| 86 | +Next, we partially fine-tune the full model starting from the best checkpoint from the first phase. |
| 87 | + |
| 88 | +```bash |
| 89 | +# Path to checkpoint |
| 90 | +ckpt_run="PHASE1_RUN_NAME" |
| 91 | +ckpt="results/algonauts23-group-encoder/${ckpt_run}/checkpoints/ckpt-best.pt" |
| 92 | + |
| 93 | +python scripts/train_group_encoder.py config/phase2_finetune.yaml \ |
| 94 | + --out_dir results --ckpt $ckpt --workers 4 $opts |
| 95 | +``` |
| 96 | + |
| 97 | +## Submission |
| 98 | + |
| 99 | +Run the [`zip_submission.sh`](scripts/zip_submission.sh) to prepare a zip file for submission to the [leaderboard](https://codalab.lisn.upsaclay.fr/competitions/9304). |
| 100 | + |
| 101 | +```bash |
| 102 | +./scripts/zip_submission.sh RESULT_DIR |
| 103 | +``` |
| 104 | + |
| 105 | +## Acknowledgements |
| 106 | + |
| 107 | +This code was built with elements and inspiration from [timm](https://github.com/huggingface/pytorch-image-models). |
| 108 | + |
| 109 | +## Citation |
| 110 | +If you find this repository helpful, please consider citing: |
| 111 | + |
| 112 | +``` |
| 113 | +@article{lane2023algonauts, |
| 114 | + author = {Connor Lane and Gregory Kiar}, |
| 115 | + title = {A Parameter-efficient Multi-subject Model for Predicting fMRI Activity}, |
| 116 | + journal = {arXiv preprint arXiv:2308.02351}, |
| 117 | + year = {2023}, |
| 118 | +} |
| 119 | +``` |
0 commit comments