czlwang / brainbert

[ICLR 2023] Code for BrainBERT

BrainBERT

BrainBERT is an modeling approach for learning self-supervised representations of intracranial electrode data. See paper for details.

We provide the training pipeline below.

The trained weights have been released (see below) and pre-training data is available upon request.

Installation

Requirements:

• pytorch >= 1.12.1
• pytorch gradual warmup scheduler

pip install -r requirements.txt

Input

It is expected that the input is intracranial electrode data that has been Laplacian re-referenced.

Using BrainBERT embeddings

• pretrained weights are available here
• see notebooks/demo.ipynb for an example input and example embedding

Upstream

BrainBERT pre-training data

The data directory should be structured as:

/pretrain_data
  |_manifests
    |_manifests.tsv  <-- each line contains the path to the example and the length
  |_<subject>
    |_<trial>
      |_<example>.npy

BrainBERT pre-training

python3 run_train.py +exp=spec2vec ++exp.runner.device=cuda ++exp.runner.multi_gpu=True \
  ++exp.runner.num_workers=64 +data=masked_spec +model=masked_tf_model_large \
  +data.data=/path/to/data ++data.val_split=0.01 +task=fixed_mask_pretrain.yaml \
  +criterion=pretrain_masked_criterion +preprocessor=stft ++data.test_split=0.01 \
  ++task.freq_mask_p=0.05 ++task.time_mask_p=0.05 ++exp.runner.total_steps=500000

Example parameters:

/path/to/data = /storage/user123/self_supervised_seeg/pretrain_data/manifests

brainbert's Issues

Code for fine-tuning

Could you please provide how to do fine-tuning? Thank you!

Handling multiple electrodes during fine-tuning

Hi there, great work!

You mention in your paper that for BrainBERT pre-training you STFT/Superlet each individual electrode to get time-frequency representations. However, I wanted to know how do you handle multiple electrodes during the fine-tuning stage.

Thank You!

Direction for generation compatible data for training/fine-tuning from edf file

Hi, I'm trying to pretrain the model with custom dataset. I saw there is some code under /data for pre-processing data. Is it possble for you to provide some direction on how to process raw EDF file and other necessary meta-data into the form compatible with this project?

Thank you!

Masking code bug

In this function: https://github.com/czlwang/BrainBERT/blob/master/util/mask_utils.py#L28
Shouldn't lines 48-49 come after the for loop on line 52?

Your masked_labels and masked input don't align

"Kind Request for Full Hardware Specifications Used in Paper's Experiments to Address Reproducibility of the Paper"

My team is trying to reproduce the results of this paper. Our system has 24 GB VRAM capacity and I am trying to reproduce the results but I lately found that my system may not be appropriate for the necessary reproducibility of results. I am able to train results on small patches (~3GB when zipped) but I am not able to train it on the entire dataset as I find the training duration required to be extremely high.

I have found that if I get access to the informtion regarding the GPUs that was used for the results of this paper. It would be really helpful if we get all the hardware specifications used for your research.

Demo data format

Can you please share some demo data so we can at least try to process ours into the required format? Thanks!