CRL-Net

Causal Representation Learning for Seizure Prediction

CRL-Net is a causal representation learning method for EEG-based epileptic seizure prediction, designed to disentangle seizure-related causal features from confounders and capture temporal contextual dependencies across consecutive EEG segments.

Structure of the Repository

crl_net/
├── train.py                 # Script for supervised seizure prediction training.
├── test.py                  # Script for model evaluation.
├── train_ts_tcc_eeg.py      # Script for pretraining the temporal encoder.
├── EEG_preprocess/          # Data preprocessing scripts for the CHB-MIT and Kaggle datasets.
├── EEG_utils/               # Utility functions and helper scripts.
├── EEG_dataset/             # PyTorch Dataset definitions and data loading utilities.
├── EEG_model/               # Model architecture definitions.
├── EEG_trainer/             # Trainer class definitions and training pipeline implementations.
├── EEG_eval/                # Evaluation metrics and result analysis utilities.
└── temporal_encoders/       # Checkpoints of pretrained temporal contextual encoders.

Quick Start

1. Clone Repository

git clone https://github.com/thuzjc/CRL-Net.git
cd CRL-Net

2. Install Dependencies

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

3. Data Preprocessing

1. Please download the raw Kaggle dataset via the following link: Download Link for the raw Kaggle dataset.
2. Please download the raw CHBMIT dataset via the following link: Download Link for the raw CHB-MIT dataset.
3. Preprocess raw data for self-supervised pretraining:

cd EEG_preprocess
bash run_preprocess_chb_cont.sh
bash run_preprocess_kaggle_cont.sh

4. Preprocess raw data for supervised training in seizure prediction:

cd EEG_preprocess
bash run_preprocess_chb.sh
bash run_preprocess_kaggle.sh

4. Run Training

1. Run cross-view self-supervised pretraining to model temporal contextual dependencies:

python train_ts_tcc_eeg.py \
  --patient_ids 1 \
  --device_number 0 \
  --encoder_backbone resnet34 \
  --feature_dim 128 \
  --context_len 4 \
  --epochs 100 \
  --batch_size 16 \
  --interictal_dir /path/to/continuous_interictal_stft_root \
  --preictal_dir /path/to/continuous_preictal_stft_root \
  --save_dir ./checkpoints_ts_tcc

The resulting pretrained checkpoints will be saved to ./checkpoints_ts_tcc/.

2. Run supervised seizure prediction training:

python train.py \
  --dataset_name CHB60 \
  --model_name MONSTB_causal \
  --patient_id 1 \
  --device_number 0 \
  --step_preictal 5 \
  --ch_num 18 \
  --encoder_backbone resnet34 \
  --temp_ckpt_dir ./temporal_encoders \
  --use_temp_encoder 1 \
  --use_cfs_loss 1 \
  --use_amp False

The counterfactual samples are constructed during this supervised training process. Common parameters:

• --use_temp_encoder 0: Don't use the pretrained temporal contextual encoder for seizure prediction.
• --use_cfs_loss 0: Disable CFS module and only use BCE loss to train the backbone module.

The resulting checkpoints will be saved to ./model/{dataset_name}/{model_name}/{patient_id}/stft/.

5. Run Evaluation

python test.py \
  --dataset_name CHB60 \
  --model_name MONSTB_causal \
  --patient_id 1 \
  --device_number 0 \
  --step_preictal 5 \
  --ch_num 18 \
  --encoder_backbone resnet34 \
  --temp_ckpt_dir ./temporal_encoders \
  --use_temp_encoder 1

If the checkpoint is obtained from training with --use_temp_encoder 0, evaluation should also be performed with --use_temp_encoder 0. Optionally, --model_files can be used to specify a particular .pth filename or path.

Acknowledgements

Built upon the excellent work of:

• Swin Transformer
• MSAN

We sincerely thank the authors for their valuable contributions.