NEDS - GitLite

KO | EN

gitlite — search

#shell #ruby #tailwindcss #react #claude #graphql #command-line #bash #cross-platform #generative-ai

NEDS

(★ 19)

Official Implementation of NEDS

File Explorer

ZIP Download

- neds_schematic.png
- eids.txt
- test_eids.txt
- train_eids.txt
- create_dataset.sh
- eval.sh
- prepare_data.sh
- run_create_dataset.sh
- train.sh
- train_multi_gpu.sh
- - __init__.py
  - base.py
  - make_loader.py
- - __init__.py
  - masker.py
  - model_output.py
  - stitcher.py
- - __init__.py
  - encoder_embeddings.py
  - mm.py
  - mm_utils.py
- - dependency_links.txt
  - PKG-INFO
  - requires.txt
  - SOURCES.txt
  - top_level.txt
- - __init__.py
  - base.py
  - make.py
- - __init__.py
  - config_utils.py
  - dataset_utils.py
  - eval_utils.py
  - ibl_data_utils.py
  - preprocess_lfp.py
  - utils.py
- create_dataset.py
- eval.py
- finetune.py
- prepare_data.py
- setup.py
- train.py
env.yaml
LICENSE
README.md

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

Command Glossary

Commands referenced in this DOCs, explained below.

🔍

`conda activate`

View Details ▼

Activate a conda environment.
See also: `conda deactivate`.

conda activate myenv

Activate an existing environment named `myenv`:

conda activate {{path/to/myenv}}

Activate an existing environment located at custom path:

conda activate --stack myenv

Stack `myenv` environment on top of a previous environment making libraries/commands/variables from both accessible:

🔍

`conda env`

View Details ▼

Manage conda environments.

conda env create {{[-f|--file]}} {{path/to/file}}

Create an environment from an environment file (YAML, TXT, etc.):

conda env remove {{[-n|--name]}} {{environment_name}}

Delete an environment and everything in it:

conda env update {{[-f|--file]}} {{path/to/file}} --prune

Update an environment based on an environment file:

🔍

`sbatch`

View Details ▼

Submit a batch job to the SLURM scheduler.

sbatch {{path/to/job.sh}}

Submit a batch job:

sbatch --job-name={{myjob}} {{path/to/job.sh}}

Submit a batch job with a custom name:

sbatch --time={{00:30:00}} {{path/to/job.sh}}

Submit a batch job with a time limit of 30 minutes:

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

We introduce a multimodal, multi-task model that enables simultaneous [Neural Encoding and Decoding at Scale (NEDS)](https://arxiv.org/abs/2504.08201). Central to our approach is a novel multi-task-masking strategy, which alternates between neural, behavioral, within-modality, and cross-modality masking.

![NEDS](https://raw.githubusercontent.com/yzhang511/NEDS/main/assets/neds_schematic.png)

## Installation

```bash
conda env create -f env.yaml    # Create conda environment

conda activate neds             # Activate conda environment
```

## Datasets and Models

**NOTE**: The IBL sessions used for pre-training are in `data/train_eids.txt`, while the held-out sessions for evaluation are in `data/test_eids.txt`.
If you want to train and evaluate on a different set of sessions, you can change the `eids.txt`, `train_eids.txt` and `test_eids.txt` files.

Download and prepare the IBL dataset. Update `base_path` and `data_path` in the scripts:

```bash
sbatch prepare_data.sh 1 EID    # Download data for a session using the provided EID

sbatch prepare_data.sh 84       # Download 84 sessions from the IBL repeated-site dataset
```

To accelerate data loading during training, we pre-save the partitioned train/val/test data and load it when needed:

```bash
sbatch create_dataset.sh 1 EID    # Save the train/val/test data for a session using the provided EID

source run_create_dataset.sh      # Save the train/val/test data for each of the 10 test sessions individually

sbatch create_dataset.sh 10 none  # Save the train/val/test data for pretraining on a set of 10 sessions
```

### Train NEDS

Train NEDS from scratch on a single session using a single GPU. Update `base_path` and `data_path` in the script:

```bash
sbatch train.sh 1 EID train mm 0 0.1 False random        # Train multi-modal model

sbatch train.sh 1 EID train encoding 0 0.1 False random  # Train encoding model

sbatch train.sh 1 EID train decoding 0 0.1 False all     # Train decoding model
```

Use `Ray Tune` for hyperparameter search. Update `num_tune_sample` in the script to change the number of random models:

```bash
sbatch train.sh 1 EID train mm 0 0.1 True random
```

Pre-train NEDS on 10 sessions using multiple GPUs. Update `--nodes` and `--ntasks` in the script to change the number of GPUs used:

```bash
sbatch train_multi_gpu.sh 10 none mm 0 0.1 all   # Pre-training requires "all"
```

Fine-tune the pre-trained 10-session NEDS model on a single held-out test session:

```bash
sbatch train.sh 10 EID finetune mm 0 0.1 False random
```

### Evaluate NEDS

To evaluate NEDS on a single session:

```bash
sbatch eval.sh 1 EID train mm 0.1 random False     # Eval a model trained from scratch (no hyperparameter search)

sbatch eval.sh 10 EID finetune mm 0.1 random False # Eval a pre-trained model fine-tuned on a held-out test session (no hyperparameter search)

sbatch eval.sh 1 EID train mm 0.1 random True      # Eval a model trained from scratch (with hyperparameter search)
```
**NOTE**: If you want to evaluate a model after hyperparameter tuning, locate the best model in the training logs and keep only its folder in `/YOUR_PATH/tune/`. Alternatively, you can change `eval.py` to automatically load the best checkpoint.

## Citation
Please cite our paper if you use this code in your own work:
```
@article{zhang2025neural,
  title={Neural Encoding and Decoding at Scale},
  author={Zhang, Yizi and Wang, Yanchen and Azabou, Mehdi and Andre, Alexandre and Wang, Zixuan and Lyu, Hanrui and Laboratory, The International Brain and Dyer, Eva and Paninski, Liam and Hurwitz, Cole},
  journal={arXiv preprint arXiv:2504.08201},
  year={2025}
}
```