Keondopark

README merge

8bedc22 2 days ago

6.67 kB

---
license: mit
language:
- en
tags:
- Sleep
- SleepFoundationModel
---
# SleepMaMi
This is the code respository for **SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures** (ICML 2026)

---
![SleepMaMi Overview](assets/SleepMaMi_Overview.png)

[![Paper](https://img.shields.io/badge/paper-arXiv-red)](https://arxiv.org/abs/2602.07628)
[![License](https://img.shields.io/badge/license-MIT-green)](LICENSE)
[![Python](https://img.shields.io/badge/python-3.10%2B-brightgreen)](#installation)
[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-SleepMaMi-FFD21E)](https://huggingface.co/gundo0102/SleepMaMi)




## Prerequisites
A CUDA GPU is required for training and inference: the FlashAttention and
Mamba (`mamba-ssm`) kernels used by the encoders have no CPU implementation.
Use a GPU that FlashAttention 2 supports (Ampere or newer, e.g. A100 / RTX 30xx+).

```
conda create -n sleepmami python==3.10
conda activate sleepmami
pip install -r requirements.txt
# conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc cuda-cudart-dev libcublas-dev libcurand-dev # if necessary
pip install flash-attn==2.7.4.post1 --no-build-isolation --cache-dir /tmp/pip_cache
apt update
apt install libgl1-mesa-glx
pip install mamba-ssm==2.2.3.post2 --no-build-isolation
```


## Data preprocessing
Please prepare pretraining and finetuning data, following [README](data/README.md) in `data` directory.

## Pretraining
### Micro-Encoder pretraining
After pretraining dataset is prepared, prepare a config file(`*.yaml`) in `configs` directory.

Example config files are provided in `configs` directory.

Use the following command to start pretraining:
```
python pretrain.py --config PT_6mod_9ds_CL
```


### Macro-Encoder pretraining
Similar to Micro-Encoder pretraining, prepare a config file in `configs` directory.

Then use the following command to start Macro-Encoder pretraining
```
python pretrain_seq.py --config PTseq_6mod_9ds_CL_full_night
```


### Downstream tasks
After pretraining is completed, finetuning on downstream tasks can be done with the following command:
```
python finetune.py --config FT_6mod_CL
```
The config file should be located in `configs` directory. `ft_task` variable should be set to specify the finetuning task (e.g. sleep_staging, apnea_seg, diag, etc.)


### Pretrained checkpoints and Inference
`inference.py` loads a pretrained / finetuned checkpoint and exposes three
entry points. Signal preprocessing reuses the same `data/data_loader.py`
pipeline used during training, so inputs are identical to what the checkpoints
were trained on.

| Task | Backbone | Output | Ckpt|
|------|----------|--------|-----|
| `micro` | `ClassificationMultiConv1d` (`classification_multi`) | per-epoch Micro feature | [micro.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/micro.safetensors)
| `macro_micro` | `ClassificationMultiConv1dMamba` (`classification_multi_full_night`) | full-night Macro + per-epoch Micro features |[macro_micro.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/macro_micro.safetensors)
| `sleep_staging` | `ClassificationMultiConv1dMamba` + classification head (SHHS1) | predicted sleep stage per epoch | [sleep_staging.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/sleep_staging.safetensors)

**Input.** Provide one `.npy` per modality holding the full night as a sequence
of 30-s epochs, shape `(N, sig_len)` where `sig_len = 3000` 
(3000 for the default 30 s @ 100 Hz). Pass a directory containing `EEG.npy`,
`EOG.npy`, … via `--npy_dir`, or list paths explicitly with `--npy`.

```bash
# 1. Micro feature from the epoch-level Micro-Encoder
python inference.py --task micro \
    --checkpoint checkpoints/micro.safetensors \
    --npy_dir /path/to/night --num_epochs 20
# 2. Macro- and Micro-features from the full-night model
python inference.py --task macro_micro \
    --checkpoint checkpoints/macro_micro.safetensors \
    --npy_dir /path/to/night
# 3. Sleep staging (Macro + Micro features + classification head)
python inference.py --task sleep_staging \
    --checkpoint checkpoints/sleep_staging.safetensors \
    --npy_dir /path/to/night
# Explicit per-modality paths instead of a directory:
python inference.py --task sleep_staging \
    --npy EEG=EEG.npy EOG=EOG.npy ECG=ECG.npy RESP=RESP.npy EMG=EMG.npy SaO2=SaO2.npy
```

Notes:
- `.npy` inputs are normalized by default to match the training pipeline: SaO2
  is scaled with `(clip(x, 0.7, 1.0) - 0.7) / 0.3` (percent values are divided
  by 100 first), and other modalities are z-normalized per epoch (re-applied by
  RevIN internally). Pass `--no_normalize` if your `.npy` arrays are already
  preprocessed.
- Released checkpoints (`checkpoints/*.safetensors`) contain model weights only
  (no optimizer / scheduler state). `--checkpoint` defaults to
  `checkpoints/<task>.safetensors`; a Lightning `.ckpt` from your own
  `finetune.py` run can also be passed, or converted to the same format with
  `train_utils.export_inference_weights`.
- Inference must run on a CUDA GPU (`--device cuda`, the default when one is
  available); running on CPU fails in the FlashAttention kernel.
- `--dtype bfloat16` reproduces the `bf16-mixed` precision used in training.
- `micro` / `macro_micro` only need an encoder (pretrained) checkpoint — the
  classification head is left uninitialized and unused. `sleep_staging`
  requires a finetuned checkpoint that contains the head.
- The functions in `inference.py` (`extract_micro_features`,
  `extract_macro_micro_features`, `predict_sleep_stages`, and the
  `load_npy_epoch_inputs` / `load_npy_full_night_inputs` helpers) can be
  imported and called directly from Python.

## SDB segmentation changes (Updated at Oct-28 2026)
The ICML paper reported Macro-F1 for SleepMaMi, but F1 of sleep disoredred breathing class for baselines. To make them consistent, we updated the arXiv paper to report F1 of sleep disordered breathing class for all models. The repository is also udpated accordingly.

The updated SDB segmentation results are as following.
![SDB segnentation](assets/SDB_Table_update_261028.png)

## BibTeX
Please use the below bibtex when you reference our work!
```
@inproceedings{
park2026sleepmami,
title={SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures},
author={Park, Keondo and Na, Younghoon and Choi, You Rim and Ryu, Hyunwoo and Shin, Hyun-Woo and Kim, Hyung-Sin},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=e9hvlHlUas}
}
```

	---
	license: mit
	language:
	- en
	tags:
	- Sleep
	- SleepFoundationModel
	---
	# SleepMaMi
	This is the code respository for SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures (ICML 2026)

	---
	![SleepMaMi Overview](assets/SleepMaMi_Overview.png)

	[![Paper](https://img.shields.io/badge/paper-arXiv-red)](https://arxiv.org/abs/2602.07628)
	[![License](https://img.shields.io/badge/license-MIT-green)](LICENSE)
	[![Python](https://img.shields.io/badge/python-3.10%2B-brightgreen)](#installation)
	[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-SleepMaMi-FFD21E)](https://huggingface.co/gundo0102/SleepMaMi)




	## Prerequisites
	A CUDA GPU is required for training and inference: the FlashAttention and
	Mamba (`mamba-ssm`) kernels used by the encoders have no CPU implementation.
	Use a GPU that FlashAttention 2 supports (Ampere or newer, e.g. A100 / RTX 30xx+).

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

	# conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc cuda-cudart-dev libcublas-dev libcurand-dev # if necessary
	pip install flash-attn==2.7.4.post1 --no-build-isolation --cache-dir /tmp/pip_cache
	apt update
	apt install libgl1-mesa-glx
	pip install mamba-ssm==2.2.3.post2 --no-build-isolation
	```


	## Data preprocessing
	Please prepare pretraining and finetuning data, following [README](data/README.md) in `data` directory.

	## Pretraining
	### Micro-Encoder pretraining
	After pretraining dataset is prepared, prepare a config file(`*.yaml`) in `configs` directory.

	Example config files are provided in `configs` directory.

	Use the following command to start pretraining:
	```
	python pretrain.py --config PT_6mod_9ds_CL
	```


	### Macro-Encoder pretraining
	Similar to Micro-Encoder pretraining, prepare a config file in `configs` directory.

	Then use the following command to start Macro-Encoder pretraining
	```
	python pretrain_seq.py --config PTseq_6mod_9ds_CL_full_night
	```


	### Downstream tasks
	After pretraining is completed, finetuning on downstream tasks can be done with the following command:
	```
	python finetune.py --config FT_6mod_CL
	```
	The config file should be located in `configs` directory. `ft_task` variable should be set to specify the finetuning task (e.g. sleep_staging, apnea_seg, diag, etc.)


	### Pretrained checkpoints and Inference
	`inference.py` loads a pretrained / finetuned checkpoint and exposes three
	entry points. Signal preprocessing reuses the same `data/data_loader.py`
	pipeline used during training, so inputs are identical to what the checkpoints
	were trained on.

	\| Task \| Backbone \| Output \| Ckpt\|
	\|------\|----------\|--------\|-----\|
	\| `micro` \| `ClassificationMultiConv1d` (`classification_multi`) \| per-epoch Micro feature \| [micro.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/micro.safetensors)
	\| `macro_micro` \| `ClassificationMultiConv1dMamba` (`classification_multi_full_night`) \| full-night Macro + per-epoch Micro features \|[macro_micro.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/macro_micro.safetensors)
	\| `sleep_staging` \| `ClassificationMultiConv1dMamba` + classification head (SHHS1) \| predicted sleep stage per epoch \| [sleep_staging.safetensors](https://huggingface.co/gundo0102/SleepMaMi/resolve/main/checkpoints/sleep_staging.safetensors)

	Input. Provide one `.npy` per modality holding the full night as a sequence
	of 30-s epochs, shape `(N, sig_len)` where `sig_len = 3000`
	(3000 for the default 30 s @ 100 Hz). Pass a directory containing `EEG.npy`,
	`EOG.npy`, … via `--npy_dir`, or list paths explicitly with `--npy`.

	```bash
	# 1. Micro feature from the epoch-level Micro-Encoder
	python inference.py --task micro \
	--checkpoint checkpoints/micro.safetensors \
	--npy_dir /path/to/night --num_epochs 20

	# 2. Macro- and Micro-features from the full-night model
	python inference.py --task macro_micro \
	--checkpoint checkpoints/macro_micro.safetensors \
	--npy_dir /path/to/night

	# 3. Sleep staging (Macro + Micro features + classification head)
	python inference.py --task sleep_staging \
	--checkpoint checkpoints/sleep_staging.safetensors \
	--npy_dir /path/to/night

	# Explicit per-modality paths instead of a directory:
	python inference.py --task sleep_staging \
	--npy EEG=EEG.npy EOG=EOG.npy ECG=ECG.npy RESP=RESP.npy EMG=EMG.npy SaO2=SaO2.npy
	```

	Notes:
	- `.npy` inputs are normalized by default to match the training pipeline: SaO2
	is scaled with `(clip(x, 0.7, 1.0) - 0.7) / 0.3` (percent values are divided
	by 100 first), and other modalities are z-normalized per epoch (re-applied by
	RevIN internally). Pass `--no_normalize` if your `.npy` arrays are already
	preprocessed.
	- Released checkpoints (`checkpoints/*.safetensors`) contain model weights only
	(no optimizer / scheduler state). `--checkpoint` defaults to
	`checkpoints/<task>.safetensors`; a Lightning `.ckpt` from your own
	`finetune.py` run can also be passed, or converted to the same format with
	`train_utils.export_inference_weights`.
	- Inference must run on a CUDA GPU (`--device cuda`, the default when one is
	available); running on CPU fails in the FlashAttention kernel.
	- `--dtype bfloat16` reproduces the `bf16-mixed` precision used in training.
	- `micro` / `macro_micro` only need an encoder (pretrained) checkpoint — the
	classification head is left uninitialized and unused. `sleep_staging`
	requires a finetuned checkpoint that contains the head.
	- The functions in `inference.py` (`extract_micro_features`,
	`extract_macro_micro_features`, `predict_sleep_stages`, and the
	`load_npy_epoch_inputs` / `load_npy_full_night_inputs` helpers) can be
	imported and called directly from Python.

	## SDB segmentation changes (Updated at Oct-28 2026)
	The ICML paper reported Macro-F1 for SleepMaMi, but F1 of sleep disoredred breathing class for baselines. To make them consistent, we updated the arXiv paper to report F1 of sleep disordered breathing class for all models. The repository is also udpated accordingly.

	The updated SDB segmentation results are as following.
	![SDB segnentation](assets/SDB_Table_update_261028.png)

	## BibTeX
	Please use the below bibtex when you reference our work!
	```
	@inproceedings{
	park2026sleepmami,
	title={SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures},
	author={Park, Keondo and Na, Younghoon and Choi, You Rim and Ryu, Hyunwoo and Shin, Hyun-Woo and Kim, Hyung-Sin},
	booktitle={Forty-third International Conference on Machine Learning},
	year={2026},
	url={https://openreview.net/forum?id=e9hvlHlUas}
	}
	```