ViRe Retrieval

ViRe is a PyTorch Lightning project for EEG-to-image feature retrieval. The currently documented pipelines are:

• src/retrieval: standard EEG-to-image feature retrieval.
• src/retrieval_GX: GX retrieval variant with an additional fusion/head configuration.

Image reconstruction pipelines are still being cleaned up and will be documented later. See README_zh.md for the Chinese version.

Project Layout

ViRe/
  configs/
    baseline_ubp.yaml       # Default retrieval config
    baseline_ubp_GX.yaml    # Default GX retrieval config
  data/
    EEG_ViRe_1202_cov_norm/ # Expected EEG tensors
    Image_feature/          # Cached visual features
    images_set/             # Images used to build feature cache
  models/                   # Shared model modules used by backbones
  src/
    retrieval/
      train.py
      data_eeg.py
      eeg_backbone.py
      utils.py
    retrieval_GX/
      train.py
      data_eeg.py
      eeg_backbone.py
      utils.py

Environment

Create a Python environment and install dependencies:

conda create -n vire python=3.10
conda activate vire
pip install -r requirements.txt

Install the PyTorch build that matches your CUDA version if the default install is not suitable for your GPU.

Required Files

Both retrieval pipelines expect data and pretrained vision weights under the repository root:

ViRe/
  data/
    EEG_ViRe_1202_cov_norm/
      sub-01/
        train.pt
        test.pt
      ...
    Image_feature/
    images_set/
  vision_backbone/
    ViT-H-14/
      open_clip_model.safetensors

If your paths differ, update self.data_dir in the relevant data_eeg.py file and the vision backbone path used when image features are generated.

Standard Retrieval

Default config: configs/baseline_ubp.yaml

Current defaults:

eeg_backbone: Ours
vision_backbone: ViT-H-14
seed: 0
alpha: 0.05
timesteps: [25, 250]
data:
  subjects: [sub-01]
models:
  brain:
    params:
      z_dim: 1024
train:
  scheduler: cosine
  min_lr: 1e-6

Run from the repository root:

python src/retrieval/train.py

Override common parameters:

python src/retrieval/train.py --eeg-backbone OursResidualBranchV2 --seed 1 --subject sub-02

python src/retrieval/train.py --start-time 25 --end-time 250 --alpha 0.05

Supported overrides:

--config
--eeg-backbone
--vision-backbone
--z-dim
--subject
--sub
--seed
--alpha
--start-time
--end-time
--info

GX Retrieval

Default config: configs/baseline_ubp_GX.yaml

Current defaults:

eeg_backbone: Ours_GX
vision_backbone: ViT-H-14
seed: 0
timesteps: [0, 300]
data:
  subjects: [sub-01]
models:
  brain:
    params:
      z_dim: 1024
      nhead: [1024, 1, 0, 0]
train:
  scheduler: cosine
  min_lr: 1e-6

Run from the repository root:

python src/retrieval_GX/train.py

Override common parameters:

python src/retrieval_GX/train.py --seed 1 --subject sub-02 --nhead 1024 1 0 0

python src/retrieval_GX/train.py --start-time 0 --end-time 300 --save-dir exp_gx

Supported overrides:

--config
--eeg-backbone
--vision-backbone
--z-dim
--subject
--sub
--seed
--start-time
--end-time
--nhead N0 N1 N2 N3
--info
--save-dir

Multiple Seeds

PowerShell example for seeds 0 through 5:

foreach ($seed in 0,1,2,3,4,5) {
    & E:\DeepLearning\miniconda3\envs\brain\python.exe src/retrieval/train.py --seed $seed --info "-ubp0.05-[25,250]-seed$seed"
}

GX example:

foreach ($seed in 0,1,2,3,4,5) {
    & E:\DeepLearning\miniconda3\envs\brain\python.exe src/retrieval_GX/train.py --seed $seed --info "-ubp-gx-seed$seed"
}

Outputs

Training writes logs, checkpoints, copied configs, and test metrics under save_dir, which defaults to exp/:

exp/
  <experiment_name>/
    <subject>/seed<seed>/
      checkpoints/
      baseline_ubp.yaml
      hparams.yaml
      test_results.json
      events.out.tfevents...

Launch TensorBoard from the repository root:

tensorboard --logdir exp

Local Notes

• The retrieval scripts currently assume CUDA training.
• Windows local runs should avoid server-style DataLoader settings such as very large num_workers.
• The warning triton not found; flop counting will not work for triton kernels is usually harmless.
• train.py resumes from ckpt_path='last' by default, so change the experiment name/info or remove the old checkpoint if you want a fresh run.