Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Health Intelligence Lab @ UCLA (https://github.com/yang-ai-lab)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,256 @@
----
-license: mit
----
+---
+license: mit
+tags:
+  - sleep
+  - eeg
+  - polysomnography
+  - foundation-model
+  - self-supervised
+  - vit
+  - biosignals
+pipeline_tag: feature-extraction
+library_name: pytorch
+language:
+  - en
+---
+
+# OSF: On Pre-training and Scaling of Sleep Foundation Models
+
+[![Paper](https://img.shields.io/badge/paper-arXiv-red)](#citation)
+[![Webpage](https://img.shields.io/badge/website-demo-blue)](https://yang-ai-lab.github.io/osf/)
+[![License](https://img.shields.io/badge/license-MIT-green)](LICENSE)
+[![Python](https://img.shields.io/badge/python-3.10%2B-brightgreen)](#installation)
+
+
+## 🔥 News
+
+- [2026-2-24] Our codebase and checkpoint is released. Full codebase for benchmarking will be public available after acceptance.
+- [2026-2-22] Our paper is out.
+
+## 📖 Introduction
+
+Polysomnography (PSG) provides the gold standard for sleep assessment but suffers from substantial heterogeneity across recording devices and cohorts.
+There have been growing efforts to build general-purpose foundation models (FMs) for sleep physiology, but lack an in-depth understanding of the pre-training process and scaling patterns that lead to more generalizable sleep FMs.
+To fill this gap, we curate a massive corpus of 166,500 hours of sleep recordings from nine public sources and establish SleepBench, a comprehensive, fully open-source benchmark.
+Leveraging SleepBench, we systematically evaluate four families of self-supervised pre-training objectives and uncover three critical findings:
+(1) existing FMs fail to generalize to missing channels at inference;
+(2) channel-invariant feature learning is essential for pre-training;
+and (3) scaling sample size, model capacity, and multi-source data mixture consistently improves downstream performance.
+With an enhanced pre-training and scaling recipe, we introduce OSF, a family of sleep FMs that achieves state-of-the-art performance across nine datasets on diverse sleep and disease prediction tasks.
+Further analysis of OSF also reveals intriguing properties in sample efficiency, hierarchical aggregation, and cross-dataset scaling.
+
+
+## 📖 Table of Contents
+
+1. [Installation](#-installation)
+2. [Quick Start](#-quick-start)
+3. [Pretrained Weights](#-pretrained-weights)
+4. [Usage](#-usage)
+5. [Benchmark Evaluations](#-benchmark-evaluations)
+6. [Supported Datasets](#-supported-datasets)
+7. [Citation](#-citation)
+
+## 💿 Installation
+
+```bash
+git clone https://huggingface.co/yang-ai-lab/OSF-Base
+cd OSF-Base
+conda env create -f environment.yml
+conda activate myenv
+```
+
+
+### Dependencies
+
+- Python >= 3.10
+- PyTorch >= 2.9.0
+- PyTorch Lightning >= 2.5.5
+
+
+## 🚀 Quick Start
+
+We provide a demo notebook (`demo.ipynb`) demonstrating how to extract embeddings from PSG signals using the pretrained model.
+
+```python
+import torch
+from osf.backbone.vit1d_cls import vit_base
+
+# Load pretrained weights (included in this repo)
+payload = torch.load("osf_backbone.pth", map_location="cpu", weights_only=False)
+meta = payload["metadata"]
+
+# Initialize model
+backbone = vit_base(
+    num_leads=meta["num_leads"],        # 12 channels
+    seq_len=meta["seq_len"],            # 1920 (64 Hz × 30 s)
+    patch_size=meta["patch_size_time"],
+    lead_wise=meta["lead_wise"],
+    patch_size_ch=meta["patch_size_ch"],
+)
+backbone.load_state_dict(payload["state_dict"])
+backbone.eval()
+
+# Extract embeddings
+# x: [B, 12, 1920] - 12-channel PSG, 64 Hz × 30 seconds
+with torch.no_grad():
+    cls_embs, patch_embs = backbone.forward_encoding(x, return_sequence=False)
+# cls_embs: [B, 768] - Global epoch-level representation
+# patch_embs: [B, 90, 768] - Local patch representations
+```
+
+## 📦 Pretrained Weights
+
+| Model | Backbone | Channels |
+|-------|----------|----------|
+| OSF | ViT-Base | 12-ch |
+
+The pretrained weights are included in this repository. You can download them via the Hugging Face Hub:
+
+```python
+from huggingface_hub import hf_hub_download
+checkpoint_path = hf_hub_download(repo_id="yang-ai-lab/OSF-Base", filename="osf_backbone.pth")
+```
+
+Or via the CLI:
+
+```bash
+huggingface-cli download yang-ai-lab/OSF-Base osf_backbone.pth
+```
+
+## 👩‍💻 Usage
+
+### Input Format
+
+Expected input format:
+- **12 PSG Channels**: ECG, EMG_Chin, EMG_LLeg, EMG_RLeg, ABD, THX, NP, SN, EOG_E1_A2, EOG_E2_A1, EEG_C3_A2, EEG_C4_A1
+- **Sample Rate**: 64 Hz
+- **Epoch Length**: 30 seconds
+- **Input Shape**: `[B, 12, 1920]`
+
+### Pretraining
+
+We support multiple self-supervised pretraining methods, for example, to launch pre-training of our OSF method, run pretraining:
+
+```bash
+python main_pretrain.py \
+    --model_name "dino_ours" \
+    --psg_encoder_name "vit_base" \
+    --batch_size 256 \
+    --lr 5e-5 \
+    --max_epochs 30 \
+    --num_devices 4 \
+    --patch_size_time 64 \
+    --patch_size_ch 4 \
+    --precision "bf16-mixed"
+```
+
+See `main_pipleines/main_pretrain.py` for more detailed settings.
+
+### Fine-tuning
+
+Fine-tune the pretrained model on downstream tasks:
+
+```bash
+python main_finetune.py \
+    --model_name "dino_ours" \
+    --ckpt_path "/path/to/pretrained/checkpoint.ckpt" \
+    --downstream_dataset_name "shhs" \
+    --eval_label "Stage" \
+    --train_data_pct 1.0 \
+    --max_steps 500 \
+    --lr 0.1 \
+    --num_devices 4
+```
+
+
+## 📊 Benchmark Evaluations
+
+### Benchmarked SSL Methods
+
+| Method | Type | Original Paper |
+|--------|------|-------------|
+| SleepFM | Contrastive | [Leave-one-out multi-modal contrastive learning](https://www.nature.com/articles/s41591-025-04133-4.pdf) |
+| SimCLR | Contrastive | [Simple Constrastive Learning](https://proceedings.mlr.press/v119/chen20j/chen20j.pdf) |
+| DINO | Self-distillation | [DINO](https://arxiv.org/pdf/2304.07193) |
+| VQ-VAE | Reconstruction | [Vector-quantized variational autoencoder](https://proceedings.neurips.cc/paper/2017/file/7a98af17e63a0ac09ce2e96d03992fbc-Paper.pdf) |
+| MAE | Reconstruction | [Masked Autoencoding](https://openaccess.thecvf.com/content/CVPR2022/papers/He_Masked_Autoencoders_Are_Scalable_Vision_Learners_CVPR_2022_paper.pdf) |
+| AR | Autoregressive | [Autoregressive Next-Token prediction](https://storage.prod.researchhub.com/uploads/papers/2020/06/01/language-models.pdf) |
+| OSF | Self-distillation | ours |
+
+### Downstream Tasks
+
+**Epoch-level Classification Tasks:**
+
+| Task | Classes | Description |
+|------|---------|-------------|
+| Sleep Stage | 4 | Awake, Light Sleep, Deep Sleep, REM classification |
+| Arousal | 2 | Arousal event detection |
+| Hypopnea | 2 | Hypopnea event detection |
+| Oxygen Desaturation | 2 | Oxygen desaturation detection |
+
+
+### Evaluation Settings
+
+| Setting | Description |
+|---------|-------------|
+| Linear Probing | Freeze backbone, train linear classifier |
+| Full Fine-tuning | Fine-tune entire model end-to-end |
+| Few-shot (k-shot) | Train with limited labeled samples |
+
+For example scripts, see `main_pipelines` and `bash_scripts` folders.
+
+## 📊 Supported Datasets
+
+We aggregated nine large-scale datasets from the National Sleep Research Resource platform.
+
+| Dataset | Full Name | Source |
+|---------|-----------|--------|
+| SHHS | Sleep Heart Health Study | NSRR |
+| CHAT | Childhood Adenotonsillectomy Trial | NSRR |
+| MROS | MrOS Sleep Study | NSRR |
+| CCSHS | Cleveland Children's Sleep and Health Study | NSRR |
+| CFS | Cleveland Family Study | NSRR |
+| MESA | Multi-Ethnic Study of Atherosclerosis | NSRR |
+| SOF | Study of Osteoporotic Fractures | NSRR |
+| WSC | Wisconsin Sleep Cohort | NSRR |
+| STAGES | Stanford Technology Analytics and Genomics in Sleep | NSRR |
+| NCHSDB | NCH Sleep DataBank | NSRR |
+
+For new users, please apply for an account and access to each of these datasets following instructions here [NSRR Registration](https://sleepdata.org/join)
+
+## 📁 Project Structure
+
+```
+OSF-Open-Sleep-Foundation-Model/
+├── osf/
+│   ├── backbone/          # ViT backbone implementations
+│   │   └── vit1d_cls.py
+│   ├── models/            # SSL model implementations
+│   │   └── dino_model_cls.py
+│   │   
+│   ├── datasets/          # Data loading utilities
+│   └── utils/             # Helper functions
+├── main_pipelines/        # Training scripts
+│   ├── main_pretrain.py
+│   └── ...
+├── bash_scripts/          # Example bash scripts
+├── osf_backbone.pth       # Pretrained model weights
+├── demo.ipynb             # Quick start demo
+├── config.py              # Dataset and channel configurations
+└── train_config.py        # Training configurations
+```
+
+
+## 📝 Citation
+
+If you use this code or models in your research, please cite our paper:
+
+```bibtex
+@article{shuai2026osf,
+  title={OSF: On Pre-training and Scaling of Sleep Foundation Models},
+  author={Shuai, Zitao and Xu, Zongzhe and Yang, David and Wang, Wei and Yang, Yuzhe},
+  journal={arXiv preprint},
+  year={2026}
+}
+```
+

config.py

@@ -0,0 +1,231 @@
+"""
+Configuration constants for sleep data processing.
+Contains dataset names, paths, channel definitions, and event labels.
+"""
+import pandas as pd
+import numpy as np
+
+# =============================================================================
+# Dataset name constants
+# =============================================================================
+SHHS = 'shhs'
+CHAT = 'chat'
+MROS = 'mros'
+CCSHS = 'ccshs'
+CFS = 'cfs'
+MESA = 'mesa'
+SOF = 'sof'
+WSC = 'wsc'
+HSP = 'hsp'
+NCHSDB = 'nchsdb'
+STAGES = 'stages'
+PATS = 'pats'
+SHHS2 = 'shhs2'
+NUMOM2B = 'numom2b'
+
+# =============================================================================
+# Data paths
+# =============================================================================
+META_PATH = '/path/to/your/nsrr/data'
+
+MASTER_SHHS = [META_PATH + "/" + SHHS + "/datasets/shhs-harmonized-dataset-0.21.0.csv"]
+MASTER_CHAT = [META_PATH + "/" + CHAT + "/datasets/chat-harmonized-dataset-0.14.0.csv"]
+MASTER_MROS = [META_PATH + "/" + MROS + "/datasets/mros-visit1-harmonized-0.6.0.csv"]
+MASTER_CCSHS = [META_PATH + "/" + CCSHS + "/datasets/ccshs-trec-harmonized-0.8.0.csv"]
+MASTER_CFS = [META_PATH + "/" + CFS + "/datasets/cfs-visit5-harmonized-dataset-0.7.0.csv"]
+MASTER_MESA = [META_PATH + "/" + MESA + "/datasets/mesa-sleep-harmonized-dataset-0.7.0.csv"]
+MASTER_SOF = [META_PATH + "/" + SOF + "/datasets/sof-visit-8-harmonized-dataset-0.8.0.csv"]
+MASTER_WSC = [META_PATH + "/" + WSC + "/datasets/wsc-harmonized-dataset-0.7.0.csv"]
+MASTER_HSP = [
+    META_PATH + "/" + HSP + "/psg-metadata/I0001_psg_metadata_2025-05-06.csv",
+    META_PATH + "/" + HSP + "/psg-metadata/I0002_psg_metadata_2025-05-06.csv",
+    META_PATH + "/" + HSP + "/psg-metadata/I0003_psg_metadata_2025-05-06.csv",
+    META_PATH + "/" + HSP + "/psg-metadata/I0004_psg_metadata_2025-05-06.csv",
+    META_PATH + "/" + HSP + "/psg-metadata/I0006_psg_metadata_2025-05-06.csv",
+]
+MASTER_STAGES = [META_PATH + "/" + STAGES + "/metadata/stages-harmonized-dataset-0.3.0.csv"]
+MASTER_NCHSDB = [META_PATH + "/" + NCHSDB + "/datasets/nchsdb-dataset-harmonized-0.3.0.csv"]
+MASTER_PATS = [META_PATH + "/" + PATS + "/datasets/pats-harmonized-dataset-0.1.0.csv"]
+
+MASTER_CSV_LIST = {
+    'shhs': MASTER_SHHS,
+    'chat': MASTER_CHAT,
+    'mros': MASTER_MROS,
+    'ccshs': MASTER_CCSHS,
+    'cfs': MASTER_CFS,
+    'mesa': MASTER_MESA,
+    'sof': MASTER_SOF,
+    'wsc': MASTER_WSC,
+    'hsp': MASTER_HSP,
+    'stages': MASTER_STAGES,
+    'pats': MASTER_PATS,
+    'nchsdb': MASTER_NCHSDB,
+}
+
+# =============================================================================
+# Channel name constants
+# =============================================================================
+# ECG channels
+ECG = 'ECG'
+ECG1 = 'ECG1'
+ECG2 = 'ECG2'
+ECG3 = 'ECG3'
+HR = 'HR'
+PPG = 'PPG'
+
+# Respiratory channels
+SPO2 = 'SPO2'
+OX = 'OX'
+ABD = 'ABD'
+THX = 'THX'
+AF = 'AF'
+NP = 'NP'
+SN = 'SN'
+
+# EOG channels
+EOG_L = 'EOG_L'
+EOG_R = 'EOG_R'
+EOG_E1_A2 = 'EOG_E1_A2'
+EOG_E2_A1 = 'EOG_E2_A1'
+
+# EMG Leg channels
+EMG_LLeg = 'EMG_LLeg'
+EMG_RLeg = 'EMG_RLeg'
+EMG_LLeg1 = 'EMG_LLeg1'
+EMG_LLeg2 = 'EMG_LLeg2'
+EMG_RLeg1 = 'EMG_RLeg1'
+EMG_RLeg2 = 'EMG_RLeg2'
+EMG_Leg = 'EMG_Leg'
+
+# Sensor Leg channels
+SENSOR_Leg = 'SENSOR_Leg'
+SENSOR_LLeg = 'SENSOR_LLeg'
+SENSOR_LLeg1 = 'SENSOR_LLeg1'
+SENSOR_LLeg2 = 'SENSOR_LLeg2'
+SENSOR_RLeg = 'SENSOR_RLeg'
+SENSOR_RLeg1 = 'SENSOR_RLeg1'
+SENSOR_RLeg2 = 'SENSOR_RLeg2'
+
+# EMG Chin channels
+EMG_Chin = 'EMG_Chin'
+EMG_RChin = 'EMG_RChin'
+EMG_LChin = 'EMG_LChin'
+EMG_CChin = 'EMG_CChin'
+
+# EEG channels (unipolar)
+EEG_C3 = 'EEG_C3'
+EEG_C4 = 'EEG_C4'
+EEG_A1 = 'EEG_A1'
+EEG_A2 = 'EEG_A2'
+EEG_O1 = 'EEG_O1'
+EEG_O2 = 'EEG_O2'
+EEG_F3 = 'EEG_F3'
+EEG_F4 = 'EEG_F4'
+
+# EEG channels (bipolar/referenced)
+EEG_C3_A2 = 'EEG_C3_A2'
+EEG_C4_A1 = 'EEG_C4_A1'
+EEG_F3_A2 = 'EEG_F3_A2'
+EEG_F4_A1 = 'EEG_F4_A1'
+EEG_O1_A2 = 'EEG_O1_A2'
+EEG_O2_A1 = 'EEG_O2_A1'
+
+# Other channels
+FPZ = 'FPZ'
+GROUND = 'GROUND'
+POS = 'POS'
+
+# =============================================================================
+# Sampling frequencies (Hz)
+# =============================================================================
+FREQ_ECG = 128
+FREQ_ECG1 = 128
+FREQ_ECG2 = 128
+FREQ_ECG3 = 128
+FREQ_HR = 1
+FREQ_PPG = 128
+
+FREQ_SPO2 = 1
+FREQ_OX = 1
+FREQ_ABD = 8
+FREQ_THX = 8
+FREQ_AF = 8
+FREQ_NP = 8
+FREQ_SN = 32
+
+FREQ_EOG_L = 64
+FREQ_EOG_R = 64
+FREQ_EOG_E1_A2 = 64
+FREQ_EOG_E2_A1 = 64
+
+FREQ_EMG_Leg = 64
+FREQ_EMG_LLeg = 64
+FREQ_EMG_RLeg = 64
+FREQ_EMG_LLeg1 = 64
+FREQ_EMG_LLeg2 = 64
+FREQ_EMG_RLeg1 = 64
+FREQ_EMG_RLeg2 = 64
+
+FREQ_SENSOR_Leg = 64
+FREQ_SENSOR_LLeg = 64
+FREQ_SENSOR_LLeg1 = 64
+FREQ_SENSOR_LLeg2 = 64
+FREQ_SENSOR_RLeg = 64
+FREQ_SENSOR_RLeg1 = 64
+FREQ_SENSOR_RLeg2 = 64
+
+FREQ_EMG_Chin = 64
+FREQ_EMG_LChin = 64
+FREQ_EMG_RChin = 64
+FREQ_EMG_CChin = 64
+
+FREQ_EEG_C3 = 64
+FREQ_EEG_C4 = 64
+FREQ_EEG_A1 = 64
+FREQ_EEG_A2 = 64
+FREQ_EEG_O1 = 64
+FREQ_EEG_O2 = 64
+FREQ_EEG_F3 = 64
+FREQ_EEG_F4 = 64
+
+FREQ_EEG_C3_A2 = 64
+FREQ_EEG_C4_A1 = 64
+FREQ_EEG_F3_A2 = 64
+FREQ_EEG_F4_A1 = 64
+FREQ_EEG_O1_A2 = 64
+FREQ_EEG_O2_A1 = 64
+
+FREQ_POS = 1
+
+# =============================================================================
+# Event annotation column names
+# =============================================================================
+EVENT_NAME_COLUMN = 'EVENT'
+START_TIME_COLUMN = 'START_SEC'
+END_TIME_COLUMN = 'END_SEC'
+
+# =============================================================================
+# Respiratory event names
+# =============================================================================
+RESPIRATORY_EVENT_CENTRAL_APNEA = 'Central Apnea'
+RESPIRATORY_EVENT_OBSTRUCTIVE_APNEA = 'Obstructive Apnea'
+RESPIRATORY_EVENT_MIXED_APNEA = 'Mixed Apnea'
+RESPIRATORY_EVENT_HYPOPNEA = 'Hypopnea'
+RESPIRATORY_EVENT_DESATURATION = 'Oxygen Desaturation'
+
+# =============================================================================
+# Limb movement event names
+# =============================================================================
+LIMB_MOVEMENT_ISOLATED = 'Limb Movement Isolated'
+LIMB_MOVEMENT_PERIODIC = 'Limb Movement Periodic'
+LIMB_MOVEMENT_ISOLATED_LEFT = 'Left Limb Movement Isolated'
+LIMB_MOVEMENT_ISOLATED_RIGHT = 'Right Limb Movement Isolated'
+LIMB_MOVEMENT_PERIODIC_LEFT = 'Left Limb Movement Periodic'
+LIMB_MOVEMENT_PERIODIC_RIGHT = 'Right Limb Movement Periodic'
+
+# =============================================================================
+# Arousal event names
+# =============================================================================
+AROUSAL_EVENT_CLASSIC = 'Arousal'
+AROUSAL_EVENT_RESPIRATORY = 'RERA'
+AROUSAL_EVENT_EMG = 'EMG-Related Arousal'

environment.yml

@@ -0,0 +1,187 @@
+name: myenv
+channels:
+  - conda-forge
+dependencies:
+  - _openmp_mutex=4.5=2_gnu
+  - bzip2=1.0.8=h4777abc_8
+  - ca-certificates=2025.8.3=hbd8a1cb_0
+  - ld_impl_linux-aarch64=2.44=h5e2c951_1
+  - libexpat=2.7.1=hfae3067_0
+  - libffi=3.4.6=he21f813_1
+  - libgcc=15.1.0=he277a41_5
+  - libgcc-ng=15.1.0=he9431aa_5
+  - libgomp=15.1.0=he277a41_5
+  - liblzma=5.8.1=h86ecc28_2
+  - libnsl=2.0.1=h86ecc28_1
+  - libsqlite=3.50.4=h022381a_0
+  - libuuid=2.41.1=h3e4203c_0
+  - libxcrypt=4.4.36=h31becfc_1
+  - libzlib=1.3.1=h86ecc28_2
+  - ncurses=6.5=ha32ae93_3
+  - openssl=3.5.2=h8e36d6e_0
+  - pip=25.2=pyh8b19718_0
+  - python=3.10.18=h256493d_0_cpython
+  - readline=8.2=h8382b9d_2
+  - setuptools=80.9.0=pyhff2d567_0
+  - tk=8.6.13=noxft_h5688188_102
+  - wheel=0.45.1=pyhd8ed1ab_1
+  - pip:
+      - absl-py==2.3.1
+      - accelerate==1.2.1
+      - aiohappyeyeballs==2.6.1
+      - aiohttp==3.12.15
+      - aiosignal==1.4.0
+      - albucore==0.0.24
+      - albumentations==2.0.8
+      - annotated-types==0.7.0
+      - asttokens==3.0.0
+      - async-timeout==5.0.1
+      - attrs==25.3.0
+      - beartype==0.22.2
+      - braceexpand==0.1.7
+      - certifi==2025.10.5
+      - cffi==2.0.0
+      - charset-normalizer==3.4.4
+      - click==8.2.1
+      - coloredlogs==15.0.1
+      - comm==0.2.3
+      - contourpy==1.3.2
+      - cosine-annealing-warmup==2.0
+      - cycler==0.12.1
+      - debugpy==1.8.17
+      - decorator==5.2.1
+      - easydict==1.13
+      - einops==0.8.1
+      - ema-pytorch==0.7.7
+      - et-xmlfile==2.0.0
+      - exceptiongroup==1.3.0
+      - executing==2.2.1
+      - filelock==3.13.1
+      - flatbuffers==25.9.23
+      - fonttools==4.59.2
+      - frozenlist==1.7.0
+      - fsspec==2024.6.1
+      - gitdb==4.0.12
+      - gitpython==3.1.45
+      - grpcio==1.75.1
+      - h5py==3.14.0
+      - hf-xet==1.1.10
+      - huggingface-hub==0.35.3
+      - humanfriendly==10.0
+      - idna==3.11
+      - imageio==2.37.0
+      - importlib-metadata==8.7.0
+      - insightface==0.7.3
+      - ipdb==0.13.13
+      - ipykernel==6.30.1
+      - ipython==8.37.0
+      - jedi==0.19.2
+      - jinja2==3.1.6
+      - joblib==1.5.2
+      - jupyter-client==8.6.3
+      - jupyter-core==5.8.1
+      - kiwisolver==1.4.9
+      - kornia==0.8.1
+      - kornia-rs==0.1.9
+      - lazy-loader==0.4
+      - lightning-utilities==0.15.2
+      - llvmlite==0.46.0
+      - loguru==0.7.3
+      - markdown==3.9
+      - markupsafe==2.1.5
+      - matplotlib==3.10.6
+      - matplotlib-inline==0.1.7
+      - ml-dtypes==0.5.3
+      - mne==1.10.1
+      - mpmath==1.3.0
+      - multidict==6.6.4
+      - munch==4.0.0
+      - nest-asyncio==1.6.0
+      - networkx==3.4.2
+      - neurokit2==0.2.12
+      - ninja==1.13.0
+      - numba==0.63.1
+      - numpy==2.2.6
+      - onnx==1.19.1
+      - onnx2torch==1.5.15
+      - onnxruntime==1.23.1
+      - opencv-python==4.12.0.88
+      - opencv-python-headless==4.12.0.88
+      - openpyxl==3.1.5
+      - packaging==24.2
+      - pandas==2.3.2
+      - parso==0.8.5
+      - pexpect==4.9.0
+      - pillow==11.0.0
+      - platformdirs==4.5.0
+      - pooch==1.8.2
+      - prettytable==3.16.0
+      - prompt-toolkit==3.0.52
+      - propcache==0.3.2
+      - protobuf==6.32.1
+      - psutil==7.1.0
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.3
+      - pyarrow==21.0.0
+      - pycparser==2.23
+      - pydantic==2.11.7
+      - pydantic-core==2.33.2
+      - pygments==2.19.2
+      - pynndescent==0.5.13
+      - pyparsing==3.2.3
+      - pysam==0.23.3
+      - python-dateutil==2.9.0.post0
+      - pytorch-lightning==2.5.5
+      - pytorch-warmup==0.2.0
+      - pytz==2025.2
+      - pyyaml==6.0.3
+      - pyzmq==27.1.0
+      - regex==2025.9.1
+      - requests==2.32.5
+      - safetensors==0.6.2
+      - scikit-image==0.25.2
+      - scikit-learn==1.7.2
+      - scipy==1.15.3
+      - seaborn==0.13.2
+      - sentencepiece==0.2.1
+      - sentry-sdk==2.37.1
+      - simsimd==6.5.3
+      - six==1.17.0
+      - smmap==5.0.2
+      - soundfile==0.13.1
+      - stack-data==0.6.3
+      - stringzilla==4.2.1
+      - sympy==1.13.1
+      - tabulate==0.9.0
+      - tensorboard==2.20.0
+      - tensorboard-data-server==0.7.2
+      - tensorboardx==2.6.4
+      - threadpoolctl==3.6.0
+      - tifffile==2025.5.10
+      - timm==1.0.19
+      - tokenizers==0.22.0
+      - tomli==2.2.1
+      - torch==2.5.1
+      - torchdiffeq==0.2.5
+      - torchmetrics==1.8.2
+      - torchtools==0.3.5
+      - torchvision==0.20.1
+      - tornado==6.5.2
+      - tqdm==4.67.1
+      - traitlets==5.14.3
+      - transformers==4.56.1
+      - typing-extensions==4.15.0
+      - typing-inspection==0.4.1
+      - tzdata==2025.2
+      - umap-learn==0.5.9.post2
+      - urllib3==2.5.0
+      - vitaldb==1.5.8
+      - wandb==0.22.1
+      - warmup-scheduler==0.3
+      - wcwidth==0.2.13
+      - webdataset==1.0.2
+      - werkzeug==3.1.3
+      - wfdb==4.3.0
+      - xxhash==3.5.0
+      - yarl==1.20.1
+      - zipp==3.23.0

finetune.bash

@@ -0,0 +1,35 @@
+
+DATASETS=("shhs" "mros")
+LABELS=("Stage" "Arousal" "Hypopnea" "Oxygen Desaturation")
+
+TRAIN_PCTS=(1.0)  
+
+declare -A MODELS
+
+MODELS["dino_ours"]="osf_vit_base.ckpt|all"
+
+for model_name in "${!MODELS[@]}"; do
+
+    IFS='|' read -r ckpt_path use_backbone <<< "${MODELS[$model_name]}"
+    
+    for dataset in "${DATASETS[@]}"; do
+        for label in "${LABELS[@]}"; do
+            for pct in "${TRAIN_PCTS[@]}"; do
+                echo "===== Model: ${model_name}, Dataset: ${dataset}, Label: ${label}, Pct: ${pct} ====="
+                
+                CUDA_VISIBLE_DEVICES=0,1,2,3 python main_finetune.py \
+                    --train_data_pct ${pct} \
+                    --max_steps 500 \
+                    --use_which_backbone "${use_backbone}" \
+                    --model_name "${model_name}" \
+                    --ckpt_path "${ckpt_path}" \
+                    --lr 0.1 \
+                    --eval_label "${label}" \
+                    --num_devices 4 \
+                    --data_source both \
+                    --include_datasets "${dataset}" \
+                    --downstream_dataset_name "${dataset}"
+            done
+        done
+    done
+done

main_pipelines/main_finetune.py

@@ -0,0 +1,261 @@
+from pprint import pprint
+import os
+from argparse import ArgumentParser, Namespace
+import datetime
+from dateutil import tz
+import random
+import numpy as np
+import torch
+import warnings
+from pytorch_lightning import seed_everything, Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, EarlyStopping
+from pytorch_lightning.loggers import WandbLogger
+
+from osf.datasets.pretrain_datamodule import SleepDataModule
+from osf.models.dino_model_cls import DINOCLSModel
+from config import *
+from train_config import *
+from osf.models.ssl_finetuner import SSLFineTuner, SSLVitalSignsRegressor
+from osf.utils.results_utils import save_results_to_json
+
+warnings.filterwarnings("ignore")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision('high')
+
+
+def main(hparams: Namespace):
+    now = datetime.datetime.now(tz.tzlocal())
+    timestamp = now.strftime("%Y_%m_%d_%H_%M_%S") + f"_{now.microsecond // 1000:03d}"
+
+    if hparams.monitor_type == "main":
+        exp_name = "finetune_12ch"
+    else:
+        exp_name = f"finetune_{hparams.monitor_type}"
+
+    if hparams.finetune_backbone:
+        exp_name = f"{exp_name}_full"
+
+    if hasattr(hparams, 'n_train_samples') and hparams.n_train_samples is not None and hparams.n_train_samples > 0:
+        pct_str = f"k{hparams.n_train_samples}"
+    elif hparams.train_data_pct < 1:
+        pct_str = f"{int(hparams.train_data_pct * 100)}pct"
+    else:
+        pct_str = "full"
+    if hparams.task_type == "classification":
+        task_label = hparams.eval_label
+    elif hparams.task_type == "regression":
+        task_label = "_".join(hparams.regression_targets)
+    else:
+        raise NotImplementedError(f"Unknown task_type: {hparams.task_type}")
+    run_name = f"{task_label}_{hparams.downstream_dataset_name}_{hparams.model_name}_{pct_str}_{timestamp}"
+    
+    ckpt_dir = os.path.join(
+        CKPT_PATH, f"logs/{exp_name}/ckpts/{run_name}")
+    os.makedirs(ckpt_dir, exist_ok=True)
+
+    if hparams.task_type == "regression":
+        ckpt_monitor = "val_mae"
+        ckpt_mode = "min"
+    else:
+        ckpt_monitor = "val_auc"
+        ckpt_mode = "max"
+    
+    callbacks = [
+        LearningRateMonitor(logging_interval="step"),
+        ModelCheckpoint(monitor=ckpt_monitor, dirpath=ckpt_dir,
+                        save_last=False, mode=ckpt_mode, save_top_k=1,
+                        auto_insert_metric_name=True),
+    ]
+    if getattr(hparams, 'early_stopping', False):
+        early_stop_callback = EarlyStopping(
+            monitor=ckpt_monitor,
+            patience=getattr(hparams, 'early_stopping_patience', 10),
+            mode=ckpt_mode,
+            verbose=True,
+        )
+        callbacks.append(early_stop_callback)
+        print(f"[INFO] Early stopping enabled: monitor={ckpt_monitor}, patience={hparams.early_stopping_patience}")
+    logger_dir = os.path.join(CKPT_PATH, f"logs/{exp_name}")
+    os.makedirs(logger_dir, exist_ok=True)
+    wandb_logger = WandbLogger(
+        project=f"{exp_name}_sleepuni", save_dir=logger_dir, name=run_name)
+    trainer = Trainer(
+        max_steps=hparams.max_steps,
+        accelerator="gpu",
+        accumulate_grad_batches=hparams.accumulate_grad_batches,
+        deterministic=True,
+        devices=hparams.num_devices,
+        strategy="ddp_find_unused_parameters_true",
+        precision=hparams.precision,
+        callbacks=callbacks,
+        logger=wandb_logger
+    )
+
+    hparams.exp_log_dir = os.path.join(
+        CKPT_PATH, f"data/{run_name}/exp_logs")
+    train_edf_cols = MONITOR_TYPE_MAP.get(hparams.monitor_type, TRAIN_EDF_COLS_UNI_ENC)
+    
+    if hparams.task_type == "regression":
+        event_cols = None
+        regression_targets = hparams.regression_targets
+        print(f"[INFO] Regression task with targets: {regression_targets}")
+    else:  # classification
+        event_cols = hparams.eval_label
+        regression_targets = None
+
+    regression_filter_config = None
+    if hparams.task_type == "regression" and "SPO2" in hparams.regression_targets:
+        if hparams.filter_spo2_min is not None or hparams.filter_spo2_max is not None:
+            spo2_filter = {}
+            if hparams.filter_spo2_min is not None:
+                spo2_filter["min"] = hparams.filter_spo2_min
+            if hparams.filter_spo2_max is not None:
+                spo2_filter["max"] = hparams.filter_spo2_max
+            regression_filter_config = {"SPO2_mean": spo2_filter}
+            print(f"[INFO] Will filter SPO2_mean with: {spo2_filter}")
+    
+    datamodule = SleepDataModule(
+            is_pretrain    = 0,
+            data_pct       = hparams.train_data_pct,
+            downstream_dataset_name  = hparams.downstream_dataset_name,
+            csv_dir        = SPLIT_DATA_FOLDER,
+            train_edf_cols = train_edf_cols,   
+            event_cols     = event_cols,
+            batch_size     = hparams.batch_size,
+            num_workers    = hparams.num_workers,
+            sample_rate = hparams.sample_rate,
+            window_size = 30,
+            data_source = hparams.data_source,
+            include_datasets = hparams.include_datasets,
+            regression_targets = regression_targets,
+            regression_filter_config = regression_filter_config,
+            n_train_samples = getattr(hparams, 'n_train_samples', None),
+            val_batch_size = getattr(hparams, 'val_batch_size', None),
+            val_data_pct = getattr(hparams, 'val_data_pct', None),
+            random_seed = hparams.seed,
+        )
+    if hparams.task_type == "regression":
+        hparams.num_classes = len(hparams.regression_targets)  # output dim
+        hparams.target_names = hparams.regression_targets
+        print(f"[INFO] Regression targets: {hparams.target_names}, num_classes={hparams.num_classes}")
+    else:  # classification
+        train_dataset = datamodule.train_dataloader().dataset
+        if hasattr(train_dataset, 'dataset'):  # It's a Subset
+            hparams.num_classes = train_dataset.dataset.num_classes
+        else:
+            hparams.num_classes = train_dataset.num_classes
+        print(f"[INFO] Classification num_classes: {hparams.num_classes}")
+    hparams.training_steps_per_epoch = len(datamodule.train_dataloader()) // hparams.accumulate_grad_batches // hparams.num_devices
+
+    if hparams.max_steps > 0:
+        hparams.total_training_steps = hparams.max_steps
+    else:
+        hparams.total_training_steps = hparams.training_steps_per_epoch * hparams.max_epochs
+    
+    print(f"Total training steps: {hparams.total_training_steps}")
+    print(f"Steps per epoch: {hparams.training_steps_per_epoch}")
+
+    class_distribution = datamodule.get_class_distribution()
+    if class_distribution is not None:
+        print(f"Class distribution: {class_distribution}")
+    hparams.class_distribution = class_distribution
+    
+    # Load pretrained DINO model
+    pretrain_model = DINOCLSModel.load_from_checkpoint(hparams.ckpt_path)
+    pprint(vars(hparams))
+
+    hparams.epochs = hparams.max_epochs
+    
+    def create_finetuner(backbones, hparams, train_edf_cols=None):
+        exclude_keys = {'train_edf_cols', 'regression_targets'}
+        hparams_dict = {k: v for k, v in vars(hparams).items() if k not in exclude_keys}
+        
+        if hparams.task_type == "regression":
+            return SSLVitalSignsRegressor(backbones=backbones, **hparams_dict)
+        else:
+            return SSLFineTuner(backbones=backbones, **hparams_dict)
+
+    # Extract ViT backbone from DINO model
+    vit = pretrain_model.encoders["all"].backbone
+    hparams.in_features = vit.width
+    print(f"[INFO] Extracted ViT backbone for dino_ours, in_features={hparams.in_features}")
+    model = create_finetuner(backbones={"all": vit}, hparams=hparams, train_edf_cols=train_edf_cols)
+
+    trainer.fit(model, datamodule=datamodule)
+    trainer.test(model, datamodule=datamodule, ckpt_path="last")
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser(description="Fine-tune pretrained model for downstream tasks.")
+    parser.add_argument("--model_name", type=str, default="dino_ours")
+    parser.add_argument("--eval_label", type=str, default="Stage",
+                        )
+    parser.add_argument("--downstream_dataset_name", type=str, default="mros",
+                        )
+    parser.add_argument("--use_which_backbone", type=str, default="all",
+                        )
+    parser.add_argument("--monitor_type", type=str, default="main",
+                        choices=["main", "type3", "type4"],
+                        help="Channel configuration: main (12ch), type3 (5ch), type4 (3ch)")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--train_data_pct", type=float, default=1.)
+    parser.add_argument("--n_train_samples", type=int, default=None,
+                        help="If set, use exactly this many training samples (overrides train_data_pct for few-shot)")
+    parser.add_argument("--data_source", type=str, default="auto",
+                        choices=["auto", "pretrain", "downstream", "both"],
+                        help="Which CSV source to use: auto (default), pretrain, downstream, or both")
+    parser.add_argument("--include_datasets", type=str, nargs="*", default=None,
+                        help="Filter by dataset names, e.g., --include_datasets shhs mros")
+    parser.add_argument("--batch_size", type=int, default=800)
+    parser.add_argument("--val_batch_size", type=int, default=None,
+                        help="Batch size for val/test (defaults to batch_size if not set, useful for few-shot)")
+    parser.add_argument("--val_data_pct", type=float, default=None,
+                        help="Percentage of val data to use (0-1, useful for few-shot to speed up validation)")
+    parser.add_argument("--patch_size_time", type=int, default=64)
+    parser.add_argument("--patch_size_ch", type=int, default=4,
+                        help="Channel patch size for 2D patchify (default: 4)")
+    parser.add_argument("--num_workers", type=int, default=32)
+    parser.add_argument("--num_devices", type=int, default=1)
+    parser.add_argument("--max_epochs", type=int, default=10)
+    parser.add_argument("--max_steps", type=int, default=2500)
+    parser.add_argument("--early_stopping", action="store_true",
+                        help="Enable early stopping based on val metric (useful for few-shot)")
+    parser.add_argument("--early_stopping_patience", type=int, default=10,
+                        help="Patience for early stopping (number of val checks without improvement)")
+    parser.add_argument("--accumulate_grad_batches", type=int, default=1)
+    parser.add_argument("--ckpt_path", type=str, default="")
+    parser.add_argument("--lr", type=float, default=1e-2)
+    parser.add_argument("--num_classes", type=int, default=2)
+    parser.add_argument("--in_features", type=int, default=256)
+    parser.add_argument("--loss_type", type=str, default="ce", choices=["ce", "focal", "balanced_softmax"],
+                        help="Loss type: 'ce' (cross-entropy), 'focal' (Focal Loss), or 'balanced_softmax' (Balanced Softmax)")
+    parser.add_argument("--focal_gamma", type=float, default=1.0,
+                        help="Gamma parameter for Focal Loss (focusing parameter)")
+    parser.add_argument("--focal_alpha", type=float, default=None,
+                        help="Alpha parameter for Focal Loss (class weighting). If None, computed from class distribution.")
+    parser.add_argument("--final_lr", type=float, default=0,
+                        help="Final learning rate for cosine annealing scheduler")
+    parser.add_argument("--use_mean_pool", action="store_true",
+                        help="Use mean pooling of all patches instead of CLS token for feature extraction")
+    parser.add_argument("--task_type", type=str, default="classification",
+                        choices=["classification", "regression"],
+                        help="Task type: classification or regression")
+    parser.add_argument("--regression_targets", type=str, nargs="*", default=["HR", "SPO2"],
+                        help="Regression targets, e.g., --regression_targets HR SPO2")
+    parser.add_argument("--filter_spo2_min", type=float, default=None,
+                        help="Filter out SPO2 values below this threshold (e.g., 70). Only applies when SPO2 is a regression target.")
+    parser.add_argument("--filter_spo2_max", type=float, default=None,
+                        help="Filter out SPO2 values above this threshold (e.g., 100). Only applies when SPO2 is a regression target.")
+    parser.add_argument("--finetune_backbone", action="store_true",
+                        help="If set, finetune the entire backbone (full finetuning); otherwise linear probing only")
+    parser.add_argument("--precision", type=str, default="32-true",
+                        choices=["32-true", "16-mixed", "bf16-mixed"],
+                        help="Training precision: 32-true (full), 16-mixed (FP16), bf16-mixed (BF16)")
+    parser.add_argument("--sample_rate", type=int, default=64,
+                        help="Input sample rate in Hz (default: 64). Use 32 for half resolution.")
+    hparams = parser.parse_args()
+
+    seed_everything(hparams.seed)
+    main(hparams)

main_pipelines/main_pretrain.py

@@ -0,0 +1,196 @@
+from pprint import pprint
+import os
+from argparse import ArgumentParser, Namespace
+import datetime
+from dateutil import tz
+import random
+import numpy as np
+import torch
+import warnings
+from datetime import timedelta
+from pytorch_lightning import seed_everything, Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, Callback
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.strategies import DDPStrategy
+
+
+class DenseStepCheckpoint(Callback):
+    """Save checkpoints at specific training steps."""
+    
+    def __init__(self, dirpath: str, save_steps: list = None):
+        super().__init__()
+        self.dirpath = dirpath
+        self.save_steps = set(save_steps) if save_steps else {1, 10, 100, 1000, 10000, 100000}
+        self.saved_steps = set()
+        
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        global_step = trainer.global_step
+        if global_step in self.save_steps and global_step not in self.saved_steps:
+            ckpt_path = os.path.join(self.dirpath, f"step={global_step}.ckpt")
+            trainer.save_checkpoint(ckpt_path)
+            self.saved_steps.add(global_step)
+            if trainer.is_global_zero:
+                print(f"[DenseStepCheckpoint] Saved checkpoint at step {global_step}: {ckpt_path}")
+
+from osf.datasets.pretrain_datamodule import SleepDataModule
+from osf.models.dino_model_cls import DINOCLSModel
+from config import *
+from train_config import *
+
+warnings.filterwarnings("ignore")
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision('high')
+
+
+torch._dynamo.config.cache_size_limit = 128
+torch._dynamo.config.optimize_ddp = False
+
+
+
+def param_stats(model: torch.nn.Module, verbose: bool = False):
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    if verbose:
+        print(f"{'Name':40s} {'Shape':20s} {'#Params':>10s} {'Train?':>6s}")
+        print("-" * 80)
+        for name, p in model.named_parameters():
+            print(f"{name:40s} {str(list(p.shape)):20s} {p.numel():10d} {str(p.requires_grad):>6s}")
+        print("-" * 80)
+    print(f"Total parameters:     {total / 1e6:.3f} M ({total})")
+    print(f"  Trainable params:   {trainable / 1e6:.3f} M ({trainable})")
+    print(f"  Frozen params:      {(total-trainable) / 1e6:.3f} M ({total-trainable})")
+def main(hparams: Namespace):
+
+    now = datetime.datetime.now(tz.tzlocal())
+    extension = now.strftime("%Y_%m_%d_%H_%M_%S")
+    extension = f"final_sleep_unimodal_{hparams.model_name}_{hparams.psg_encoder_name}_bz{hparams.batch_size}_{extension}"
+    ckpt_dir = os.path.join(
+        CKPT_PATH, f"logs/sleepuni/ckpts/{extension}")
+    os.makedirs(ckpt_dir, exist_ok=True)
+    if hparams.model_name in MODEL_LIST:
+        callbacks = [
+            LearningRateMonitor(logging_interval="step"),
+            ModelCheckpoint(monitor="val/loss", dirpath=ckpt_dir,
+                            save_last=True, every_n_epochs=2, mode="min", save_top_k=-1,
+                            save_on_train_epoch_end=False, auto_insert_metric_name=True),
+        ]
+        if hparams.dense_ckpt:
+            dense_ckpt_dir = os.path.join(ckpt_dir, "dense_steps")
+            os.makedirs(dense_ckpt_dir, exist_ok=True)
+            callbacks.append(DenseStepCheckpoint(
+                dirpath=dense_ckpt_dir,
+                save_steps=hparams.dense_ckpt_steps
+            ))
+    else:
+        raise NotImplementedError
+    logger_dir = os.path.join(CKPT_PATH, "logs/sleepuni")
+    os.makedirs(logger_dir, exist_ok=True)
+    print("wandb logger dir: ", logger_dir)
+    wandb_logger = WandbLogger(
+        project=hparams.wandb_proj_name + f'final_{hparams.model_name}_{hparams.psg_encoder_name}_bz{hparams.batch_size}', save_dir=logger_dir, name=extension)
+
+    strategy = DDPStrategy(
+        find_unused_parameters=True,
+        static_graph=False,
+        timeout=timedelta(minutes=15),
+    )
+
+    trainer = Trainer(
+        max_epochs=hparams.max_epochs,
+        accelerator="gpu",
+        accumulate_grad_batches=hparams.accumulate_grad_batches,
+        devices=hparams.num_devices,
+        num_nodes=hparams.num_nodes,
+        precision=hparams.precision,
+        gradient_clip_val=3.0,
+        gradient_clip_algorithm="norm",
+        strategy=strategy,
+        callbacks=callbacks,
+        logger=wandb_logger,
+        log_every_n_steps=10,
+    )
+
+    hparams.exp_log_dir = os.path.join(
+        CKPT_PATH, f"data/{extension}/exp_logs")
+    train_edf_cols = MONITOR_TYPE_MAP.get(hparams.monitor_type, TRAIN_EDF_COLS_UNI_ENC)
+    hparams.num_leads = len(train_edf_cols)
+    
+    dm = SleepDataModule(
+            is_pretrain    = 1,
+            csv_dir        = SPLIT_DATA_FOLDER,
+            train_edf_cols = train_edf_cols,
+            batch_size     = hparams.batch_size,
+            num_workers    = hparams.num_workers,
+            data_pct       = hparams.train_data_pct,
+            window_size = 30,
+            sample_rate = 64,
+            val_dataset_list = hparams.val_dataset_list,
+            data_source = hparams.data_source,
+            include_datasets = hparams.include_datasets,
+        )
+
+    hparams.simclr_augmentation = AUGMENTATION_MAP.get(hparams.model_name, "none")
+
+    # Create DINO model
+    model = DINOCLSModel(**vars(hparams))
+    model.training_steps_per_epoch = len(dm.train_dataloader()) // hparams.accumulate_grad_batches // hparams.num_devices
+    model.teacher_temp_warmup_iters = model.training_steps_per_epoch * 0.1 * hparams.max_epochs
+    print(f"[INFO] DINO teacher warmup steps: {model.teacher_temp_warmup_iters}")
+    pprint(vars(hparams))
+
+    if hparams.ckpt_path:
+        trainer.fit(model, datamodule = dm, ckpt_path=hparams.ckpt_path)
+    else:
+        trainer.fit(model, datamodule = dm)
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser(description="Pretraining DINO model for sleep PSG data.")
+    parser.add_argument("--model_name", type=str, default="dino_ours",
+                        choices=MODEL_LIST)
+    
+    parser.add_argument("--psg_encoder_name", type=str, default="vit_base")
+    parser.add_argument("--val_dataset_list", default=PRETRAIN_VAL_DATASET_LIST)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--train_data_pct", type=float, default=1.)
+    parser.add_argument("--data_source", type=str, default="auto",
+                        choices=["auto", "pretrain", "downstream", "both"])
+    parser.add_argument("--include_datasets", type=str, nargs="*", default=None)
+    parser.add_argument("--monitor_type", type=str, default="main",
+                        choices=["main", "type3", "type4"],
+                        help="Channel configuration: main (12ch), type3 (5ch), type4 (3ch)")
+    parser.add_argument("--batch_size", type=int, default=32)
+    parser.add_argument("--patch_size_time", type=int, default=4)
+    parser.add_argument("--patch_size_ch", type=int, default=4)
+    parser.add_argument("--use_2d_pos_embed", type=bool, default=True)
+    parser.add_argument("--sample_rate", type=int, default=64)
+    parser.add_argument("--num_workers", type=int, default=64)
+    parser.add_argument("--num_devices", type=int, default=4)
+    parser.add_argument("--num_nodes", type=int, default=1)
+    parser.add_argument("--max_epochs", type=int, default=30)
+    parser.add_argument("--accumulate_grad_batches", type=int, default=1)
+    parser.add_argument("--precision", type=str, default="32-true")
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--text_encoder_name", type=str, default="google/flan-t5-base")
+    parser.add_argument("--lead_wise", type=int, default=0)
+    parser.add_argument("--use_lead_embedding", type=int, default=1)
+    # DINO-specific args
+    parser.add_argument("--koleo_lambda", type=float, default=0.0)
+    parser.add_argument("--ibot_lambda", type=float, default=0.0)
+    parser.add_argument("--dino_out_dim", type=int, default=2048)
+    parser.add_argument("--dino_patch_out_dim", type=int, default=2048)
+    parser.add_argument("--dino_hidden_dim", type=int, default=2048)
+    parser.add_argument("--dino_bottleneck_dim", type=int, default=256)
+    parser.add_argument("--wandb_proj_name", type=str, default="sleepuni")
+    parser.add_argument("--ckpt_path", type=str, default=None)
+    parser.add_argument("--dense_ckpt", action="store_true")
+    parser.add_argument("--dense_ckpt_steps", type=int, nargs="+", default=[10, 100, 200, 400, 500, 800, 1000, 1600, 2500, 3200, 6400, 10000, 12500, 12800, 25600, 51200, 62500, 100000])
+    
+
+
+    hparams = parser.parse_args()
+    
+    seed_everything(hparams.seed)
+    main(hparams)

osf/backbone/pos_embed.py

@@ -0,0 +1,76 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+import numpy as np
+
+import torch
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2)
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        new_size = int(num_patches ** 0.5)
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed

osf/backbone/vit1d.py

@@ -0,0 +1,209 @@
+"""
+1D Vision Transformer for time-series signals.
+
+Patchify modes:
+- lead_wise=0: 1D patchify (all channels in one patch), no lead embedding
+- lead_wise=1: 2D patchify (channel groups), with lead embedding by default
+"""
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class DropPath(nn.Module):
+    def __init__(self, drop_prob: float, scale_by_keep: bool = True):
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        if self.drop_prob <= 0. or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0 and self.scale_by_keep:
+            random_tensor.div_(keep_prob)
+        return x * random_tensor
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim: int, fn: nn.Module):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, hidden_dim: int, drop_out_rate=0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(drop_out_rate),
+            nn.Linear(hidden_dim, output_dim),
+            nn.Dropout(drop_out_rate)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, heads: int = 8, dim_head: int = 64,
+                 qkv_bias: bool = True, drop_out_rate: float = 0., attn_drop_out_rate: float = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == input_dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(attn_drop_out_rate)
+        self.to_qkv = nn.Linear(input_dim, inner_dim * 3, bias=qkv_bias)
+
+        if project_out:
+            self.to_out = nn.Sequential(nn.Linear(inner_dim, output_dim), nn.Dropout(drop_out_rate))
+        else:
+            self.to_out = nn.Identity()
+
+    def forward(self, x):
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, hidden_dim: int, heads: int = 8,
+                 dim_head: int = 32, qkv_bias: bool = True, drop_out_rate: float = 0.,
+                 attn_drop_out_rate: float = 0., drop_path_rate: float = 0.):
+        super().__init__()
+        attn = Attention(input_dim, output_dim, heads, dim_head, qkv_bias, drop_out_rate, attn_drop_out_rate)
+        self.attn = PreNorm(input_dim, attn)
+        self.droppath1 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+        ff = FeedForward(output_dim, output_dim, hidden_dim, drop_out_rate)
+        self.ff = PreNorm(output_dim, ff)
+        self.droppath2 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def forward(self, x):
+        x = self.droppath1(self.attn(x)) + x
+        x = self.droppath2(self.ff(x)) + x
+        return x
+
+
+class ViT(nn.Module):
+    def __init__(self,
+                 num_leads: int,
+                 seq_len: int,
+                 patch_size: int,
+                 lead_wise=0,
+                 patch_size_ch=4,
+                 use_lead_embedding: bool = True,
+                 width: int = 768,
+                 depth: int = 12,
+                 mlp_dim: int = 3072,
+                 heads: int = 12,
+                 dim_head: int = 64,
+                 qkv_bias: bool = True,
+                 drop_out_rate: float = 0.,
+                 attn_drop_out_rate: float = 0.,
+                 drop_path_rate: float = 0.,
+                 **kwargs):
+        super().__init__()
+        assert seq_len % patch_size == 0
+        num_patches = seq_len // patch_size
+        self.lead_wise = lead_wise
+        self.use_lead_embedding = use_lead_embedding
+
+        if lead_wise == 0:
+            self.to_patch_embedding = nn.Conv1d(num_leads, width, kernel_size=patch_size, stride=patch_size, bias=False)
+            self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, width))
+        else:
+            self.to_patch_embedding = nn.Conv2d(1, width, kernel_size=(patch_size_ch, patch_size),
+                                                stride=(patch_size_ch, patch_size), bias=False)
+            self.pos_embedding = nn.Parameter(torch.randn(1, num_patches * num_leads // patch_size_ch, width))
+            if use_lead_embedding:
+                self.lead_emb = nn.Embedding(num_leads // patch_size_ch, width)
+            else:
+                self.lead_emb = None
+
+        self.dropout = nn.Dropout(drop_out_rate)
+        self.depth = depth
+        self.width = width
+
+        drop_path_rate_list = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        for i in range(depth):
+            block = TransformerBlock(width, width, mlp_dim, heads, dim_head, qkv_bias,
+                                     drop_out_rate, attn_drop_out_rate, drop_path_rate_list[i])
+            self.add_module(f'block{i}', block)
+
+        self.norm = nn.LayerNorm(width)
+        self.head = nn.Identity()
+
+    def _patchify_and_embed(self, series: torch.Tensor) -> torch.Tensor:
+        """Patchify input and add positional/lead embeddings. [B,C,T] -> [B,N,D]"""
+        if self.lead_wise == 0:
+            x = self.to_patch_embedding(series)  # [B, D, N]
+            x = rearrange(x, 'b c n -> b n c')   # [B, N, D]
+            x = x + self.pos_embedding[:, :x.size(1), :].to(x.device)
+        else:
+            x = self.to_patch_embedding(series.unsqueeze(1))  # [B, D, Lr, Nt]
+            Lr, Nt = x.shape[-2], x.shape[-1]
+            x = rearrange(x, 'b c lr nt -> b (lr nt) c')      # [B, N, D]
+            x = x + self.pos_embedding[:, :x.size(1), :].to(x.device)
+            if self.use_lead_embedding and self.lead_emb is not None:
+                row_ids = torch.arange(Lr, device=x.device).repeat_interleave(Nt)
+                x = x + self.lead_emb(row_ids)[None, :, :]
+        return x
+
+    def forward_encoding(self, series: torch.Tensor) -> torch.Tensor:
+        """Encode series. Returns [B,D] (mean pooled)."""
+        x = self._patchify_and_embed(series)
+        x = self.dropout(x)
+        for i in range(self.depth):
+            x = getattr(self, f'block{i}')(x)
+        x = x.mean(dim=1)
+        return self.norm(x)
+
+    def forward(self, series):
+        x = self.forward_encoding(series)
+        return self.head(x)
+
+    def reset_head(self, num_classes=1):
+        del self.head
+        self.head = nn.Linear(self.width, num_classes)
+
+
+def vit_nano(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size,
+               width=128, depth=6, heads=4, mlp_dim=512, **kwargs)
+
+
+def vit_tiny(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size,
+               width=192, depth=12, heads=3, mlp_dim=768, **kwargs)
+
+
+def vit_small(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size,
+               width=384, depth=12, heads=6, mlp_dim=1536, **kwargs)
+
+
+def vit_middle(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size,
+               width=512, depth=12, heads=8, mlp_dim=2048, **kwargs)
+
+
+def vit_base(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(num_leads=num_leads, num_classes=num_classes, seq_len=seq_len, patch_size=patch_size,
+               width=768, depth=12, heads=12, mlp_dim=3072, **kwargs)

osf/backbone/vit1d_cls.py

@@ -0,0 +1,363 @@
+"""
+1D Vision Transformer with CLS token support.
+
+Patchify modes:
+- lead_wise=0: 1D patchify (all channels in one patch)
+- lead_wise=1: 2D patchify (channel groups)
+
+Note: lead_emb is DEPRECATED and not used in data flow. It is kept only for
+checkpoint compatibility. Do NOT add lead_emb usage without careful consideration.
+"""
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+class DropPath(nn.Module):
+    '''
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    '''
+    def __init__(self, drop_prob: float, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        if self.drop_prob <= 0. or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0 and self.scale_by_keep:
+            random_tensor.div_(keep_prob)
+        return x * random_tensor
+
+
+class PreNorm(nn.Module):
+    def __init__(self,
+                 dim: int,
+                 fn: nn.Module):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+
+class FeedForward(nn.Module):
+    """
+    MLP Module with GELU activation fn + dropout.
+    """
+    def __init__(self,
+                 input_dim: int,
+                 output_dim: int,
+                 hidden_dim: int,
+                 drop_out_rate=0.):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(input_dim, hidden_dim),
+                                 nn.GELU(),
+                                 nn.Dropout(drop_out_rate),
+                                 nn.Linear(hidden_dim, output_dim),
+                                 nn.Dropout(drop_out_rate))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 output_dim: int,
+                 heads: int = 8,
+                 dim_head: int = 64,
+                 qkv_bias: bool = True,
+                 drop_out_rate: float = 0.,
+                 attn_drop_out_rate: float = 0.):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == input_dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(attn_drop_out_rate)
+        self.to_qkv = nn.Linear(input_dim, inner_dim * 3, bias=qkv_bias)
+
+        if project_out:
+            self.to_out = nn.Sequential(nn.Linear(inner_dim, output_dim),
+                                        nn.Dropout(drop_out_rate))
+        else:
+            self.to_out = nn.Identity()
+
+    def forward(self, x):
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out = self.to_out(out)
+        return out
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 output_dim: int,
+                 hidden_dim: int,
+                 heads: int = 8,
+                 dim_head: int = 32,
+                 qkv_bias: bool = True,
+                 drop_out_rate: float = 0.,
+                 attn_drop_out_rate: float = 0.,
+                 drop_path_rate: float = 0.):
+        super().__init__()
+        attn = Attention(input_dim=input_dim,
+                         output_dim=output_dim,
+                         heads=heads,
+                         dim_head=dim_head,
+                         qkv_bias=qkv_bias,
+                         drop_out_rate=drop_out_rate,
+                         attn_drop_out_rate=attn_drop_out_rate)
+        self.attn = PreNorm(dim=input_dim,
+                            fn=attn)
+        self.droppath1 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+        ff = FeedForward(input_dim=output_dim,
+                         output_dim=output_dim,
+                         hidden_dim=hidden_dim,
+                         drop_out_rate=drop_out_rate)
+        self.ff = PreNorm(dim=output_dim,
+                          fn=ff)
+        self.droppath2 = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def forward(self, x):
+        x = self.droppath1(self.attn(x)) + x
+        x = self.droppath2(self.ff(x)) + x
+        return x
+
+
+class ViT(nn.Module):
+
+    def __init__(self,
+                 num_leads: int,
+                 seq_len: int,
+                 patch_size: int,
+                 lead_wise: int = 0,
+                 patch_size_ch: int = 4,
+                 width: int = 768,
+                 depth: int = 12,
+                 mlp_dim: int = 3072,
+                 heads: int = 12,
+                 dim_head: int = 64,
+                 qkv_bias: bool = True,
+                 drop_out_rate: float = 0.,
+                 attn_drop_out_rate: float = 0.,
+                 drop_path_rate: float = 0.,
+                 **kwargs):
+        super().__init__()
+        assert seq_len % patch_size == 0
+        num_patches_time = seq_len // patch_size
+
+        self.lead_wise = lead_wise
+        self.width = width
+        self.depth = depth
+
+        if lead_wise == 0:
+            self.to_patch_embedding = nn.Conv1d(num_leads, width, kernel_size=patch_size,
+                                                stride=patch_size, bias=False)
+            N_max = num_patches_time
+            self.lead_emb = None
+        else:
+            self.to_patch_embedding = nn.Conv2d(1, width,
+                                                kernel_size=(patch_size_ch, patch_size),
+                                                stride=(patch_size_ch, patch_size),
+                                                bias=False)
+            Lr = num_leads // patch_size_ch
+            N_max = Lr * num_patches_time
+            self.lead_emb = nn.Embedding(Lr, width)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, width))
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        self.pos_embedding = nn.Parameter(torch.zeros(1, N_max + 1, width))
+        nn.init.trunc_normal_(self.pos_embedding, std=0.02)
+
+        self.dropout = nn.Dropout(drop_out_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        for i in range(depth):
+            block = TransformerBlock(input_dim=width, output_dim=width,
+                                     hidden_dim=mlp_dim, heads=heads, dim_head=dim_head,
+                                     qkv_bias=qkv_bias, drop_out_rate=drop_out_rate,
+                                     attn_drop_out_rate=attn_drop_out_rate,
+                                     drop_path_rate=dpr[i])
+            self.add_module(f'block{i}', block)
+
+        self.norm = nn.LayerNorm(width)
+        self.head = nn.Identity()
+
+
+    def to_tokens_2d(self, series: torch.Tensor,
+                     patch_size_ch: int | None = None,
+                     patch_size_time: int | None = None):
+        """Patchify only (no pos embedding). Returns (tokens, meta)."""
+        B, L, T = series.shape
+
+        if self.lead_wise == 0:
+            x = self.to_patch_embedding(series)         # [B,C,Nt]
+            Nt = x.shape[-1]
+            x = rearrange(x, 'b c n -> b n c')          # [B,Nt,C]
+            meta = dict(lead_wise=0, L=L, Nt=Nt, pz_ch=1)
+            return x, meta
+
+        # lead_wise == 1
+        if patch_size_ch is None or patch_size_time is None:
+            kch, kt = self.to_patch_embedding.kernel_size
+            patch_size_ch = patch_size_ch or kch
+            patch_size_time = patch_size_time or kt
+        assert L % patch_size_ch == 0 and T % patch_size_time == 0
+
+        x = series.unsqueeze(1)                          # [B,1,L,T]
+        x = self.to_patch_embedding(x)                   # [B,C,Lr,Nt]
+        Lr, Nt = x.shape[-2], x.shape[-1]
+        x = rearrange(x, 'b c lr nt -> b (lr nt) c')     # [B,Lr*Nt,C]
+        meta = dict(lead_wise=1, L=L, Nt=Nt, pz_ch=patch_size_ch)
+        return x, meta
+
+    def forward_encoding(self, series: torch.Tensor,
+                         return_sequence: bool = False):
+        """Encode with CLS token. Returns (cls, patches) or full sequence if return_sequence=True."""
+        tokens, meta = self.to_tokens_2d(series)
+        B = tokens.size(0)
+        cls_tok = self.cls_token.expand(B, -1, -1)
+        x = torch.cat([cls_tok, tokens], dim=1)         # [B,N+1,C]
+
+        pe = self.pos_embedding[:, :x.size(1), :].to(x.device)
+
+        x = x + pe
+
+        x = self._run_blocks(x)
+        if return_sequence:
+            return x
+        cls, patches = x[:, 0], x[:, 1:]
+        
+        return cls, patches
+
+
+    def _run_blocks(self, x: torch.Tensor):
+        x = self.dropout(x)
+        for i in range(self.depth):
+            x = getattr(self, f'block{i}')(x)
+        x = self.norm(x)
+        return self.head(x)
+
+    def forward(self, series: torch.Tensor):
+        cls, _ = self.forward_encoding(series, return_sequence=False)
+        return cls
+
+    def forward_avg_pool(self, series: torch.Tensor):
+        """Returns avg-pooled patch embeddings. series: [B,C,T] -> [B,D]"""
+        _, patches = self.forward_encoding(series, return_sequence=False)  # [B,N,D]
+        return patches.mean(dim=1)  # [B,D]
+
+    def reset_head(self, num_classes=1):
+        del self.head
+        self.head = nn.Linear(self.width, num_classes)
+
+
+    
+def vit_nano(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    model_args = dict(num_leads=num_leads,
+                      num_classes=num_classes,
+                      seq_len=seq_len,
+                      patch_size=patch_size,
+                      width=128,
+                      depth=6,
+                      heads=4,
+                      mlp_dim=512,
+                      **kwargs)
+    return ViT(**model_args)
+
+
+def vit_tiny(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    model_args = dict(num_leads=num_leads,
+                      num_classes=num_classes,
+                      seq_len=seq_len,
+                      patch_size=patch_size,
+                      width=192,
+                      depth=12,
+                      heads=3,
+                      mlp_dim=768,
+                      **kwargs)
+    return ViT(**model_args)
+
+
+def vit_small(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    model_args = dict(num_leads=num_leads,
+                      num_classes=num_classes,
+                      seq_len=seq_len,
+                      patch_size=patch_size,
+                      width=384,
+                      depth=12,
+                      heads=6,
+                      mlp_dim=1536,
+                      **kwargs)
+    return ViT(**model_args)
+
+
+def vit_middle(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    model_args = dict(num_leads=num_leads,
+                      num_classes=num_classes,
+                      seq_len=seq_len,
+                      patch_size=patch_size,
+                      width=512,
+                      depth=12,
+                      heads=8,
+                      mlp_dim=2048,
+                      **kwargs)
+    return ViT(**model_args)
+
+
+def vit_base(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    model_args = dict(num_leads=num_leads,
+                      num_classes=num_classes,
+                      seq_len=seq_len,
+                      patch_size=patch_size,
+                      width=768,
+                      depth=12,
+                      heads=12,
+                      mlp_dim=3072,
+                      **kwargs)
+    return ViT(**model_args)
+
+
+def vit_large(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(
+        num_leads=num_leads,
+        num_classes=num_classes,
+        seq_len=seq_len,
+        patch_size=patch_size,
+        width=1024,
+        depth=24,
+        heads=16,
+        mlp_dim=4096,
+        **kwargs
+    )
+
+def vit_xl(num_leads, num_classes=1, seq_len=5000, patch_size=50, **kwargs):
+    return ViT(
+        num_leads=num_leads,
+        num_classes=num_classes,
+        seq_len=seq_len,
+        patch_size=patch_size,
+        width=1536,
+        depth=24,        
+        heads=24,
+        mlp_dim=6144,
+        **kwargs
+    )

osf/datasets/augmentations.py

@@ -0,0 +1,81 @@
+"""
+Data augmentations for SSL pretraining (SimCLR, DINO).
+"""
+import torch
+import torch.nn.functional as F
+from typing import Tuple
+
+
+@torch.no_grad()
+def random_time_crop(
+    x: torch.Tensor, 
+    ratio: Tuple[float, float] | float = (0.6, 0.9),
+    *, 
+    resize_back: bool = True,
+    align_to: int | None = 40
+) -> torch.Tensor:
+    """
+    Randomly crop a contiguous sub-sequence per sample, optionally resize back to original T.
+    
+    Args:
+        x: (B, C, T)
+        ratio: crop length ratio in [low, high] or a float
+        resize_back: if True, linearly interpolate the cropped view back to length T
+        align_to: if not None, crop length is rounded to a multiple of align_to (>= align_to)
+    """
+    assert x.dim() == 3, f"expected (B,C,T), got {tuple(x.shape)}"
+    B, C, T = x.shape
+    dev = x.device
+
+    def _sample_L() -> int:
+        if isinstance(ratio, (tuple, list)):
+            a, b = float(ratio[0]), float(ratio[1])
+            r = torch.empty((), device=dev).uniform_(a, b).item()
+        else:
+            r = float(ratio)
+        L = max(2, int(round(T * r)))
+        if align_to and align_to > 1:
+            L = max(align_to, int(round(L / align_to)) * align_to)
+        return min(L, T)
+
+    Ls = [_sample_L() for _ in range(B)]
+    outs = []
+    for b in range(B):
+        L = Ls[b]
+        max_start = max(0, T - L)
+        s = int(torch.randint(0, max_start + 1, (1,), device=dev).item())
+        v = x[b, :, s:s+L]  # (C, L)
+        if resize_back and v.shape[-1] != T:
+            v = F.interpolate(v[None], size=T, mode="linear", align_corners=False)[0]
+        outs.append(v)
+    return torch.stack(outs, dim=0)
+
+
+@torch.no_grad()
+def channel_dropout(
+    x: torch.Tensor, 
+    drop_prob: float = 0.2, 
+    min_keep: int = 1
+) -> torch.Tensor:
+    """
+    Drop entire channels to zero with probability drop_prob (per sample, per channel).
+    Ensures at least `min_keep` channels remain active in each sample.
+    
+    Args:
+        x: (B, C, T)
+        drop_prob: probability to drop each channel
+        min_keep: minimum number of channels to keep per sample
+    """
+    assert x.dim() == 3
+    B, C, T = x.shape
+    mask = (torch.rand(B, C, 1, device=x.device, dtype=x.dtype) > drop_prob).to(x.dtype)
+    
+    # Ensure at least min_keep channels kept
+    keep = mask.sum(dim=1, keepdim=True)  # (B, 1, 1)
+    need = (keep < min_keep).squeeze(-1).squeeze(-1)  # (B,)
+    if need.any():
+        for b in torch.where(need)[0]:
+            idx = torch.randperm(C, device=x.device)[:min_keep]
+            mask[b, idx, 0] = 1.0
+    
+    return x * mask

osf/datasets/pretrain_datamodule.py

@@ -0,0 +1,303 @@
+import os
+from typing import List, Sequence, Optional, Dict, Union
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import pytorch_lightning as pl
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning.utilities.types import EVAL_DATALOADERS
+from torch.utils.data import DataLoader, Subset
+
+from osf.datasets.pretrain_dataset import SleepEpochDataset
+
+
+
+class SleepDataModule(LightningDataModule):
+
+    def __init__(
+        self,
+        csv_dir: str | Path,
+        *,
+        is_pretrain,
+        data_pct = 1,
+        val_dataset_list: Optional[List[str]] = None,
+        downstream_dataset_name  = None,
+        batch_size: int = 128,
+        num_workers: int = 4,
+        patient_cols: Optional[Union[str, Sequence[str]]] = None,
+        event_cols: Optional[Union[str, Sequence[str]]] = None,
+        train_edf_cols: Sequence[str] | None,
+        transforms=None,
+        n_views: int = 1,
+        cache_size: int = 8,
+        sample_rate: int = 128,
+        window_size: int = 30,
+        pin_memory: bool = False,
+        persistent_workers: bool = False,
+        data_source: str = "auto",
+        include_datasets: Optional[List[str]] = None,
+        regression_targets: Optional[List[str]] = None,
+        regression_filter_config: Optional[Dict] = None,
+        n_train_samples: Optional[int] = None,
+        val_batch_size: Optional[int] = None,
+        val_data_pct: Optional[float] = None,
+        return_all_event_cols: bool = False,
+        return_nsrrid: bool = False,
+        random_seed: int = 42,
+    ):
+        super().__init__()
+        self.save_hyperparameters(ignore=["transforms"])  
+        self.downstream_dataset_name  = downstream_dataset_name
+        self.csv_dir   = csv_dir
+        self.transforms = transforms
+        self.n_views    = n_views
+        self.pin_memory = pin_memory
+        self.persistent_workers = persistent_workers
+        self.is_pretrain = is_pretrain
+        self.patient_cols = patient_cols
+        self.event_cols = event_cols
+        self.data_pct = data_pct
+        self.data_source = data_source
+        self.include_datasets = include_datasets
+        self.regression_targets = regression_targets
+        self.regression_filter_config = regression_filter_config
+        self.n_train_samples = n_train_samples
+        self.val_batch_size = val_batch_size
+        self.val_data_pct = val_data_pct
+        self.return_all_event_cols = return_all_event_cols
+        self.return_nsrrid = return_nsrrid
+        self.random_seed = random_seed
+
+    def train_dataloader(self):
+        if self.is_pretrain == 1:
+            train_set = SleepEpochDataset(
+                    csv_dir       = self.csv_dir,
+                    split         = "pretrain",
+                    data_pct      = self.data_pct,
+                    train_edf_cols= self.hparams.train_edf_cols,
+                    transform     = self.transforms,
+                    sample_rate   = self.hparams.sample_rate,
+                    window_size   = self.hparams.window_size,
+                    cache_size    = self.hparams.cache_size,
+                    data_source   = self.data_source,
+                    include_datasets = self.include_datasets,
+                )
+            persistent_workers = self.persistent_workers
+        else:
+            train_set = SleepEpochDataset(
+                    csv_dir       = self.csv_dir,
+                    split         = "train",
+                    data_pct      = self.data_pct,
+                    patient_cols  = self.patient_cols,
+                    event_cols    = self.event_cols,
+                    train_edf_cols= self.hparams.train_edf_cols,
+                    transform     = self.transforms,
+                    sample_rate   = self.hparams.sample_rate,
+                    window_size   = self.hparams.window_size,
+                    cache_size    = self.hparams.cache_size,
+                    downstream_dataset_name  = self.downstream_dataset_name,
+                    data_source   = self.data_source,
+                    include_datasets = self.include_datasets,
+                    regression_targets = self.regression_targets,
+                    regression_filter_config = self.regression_filter_config,
+                    return_all_event_cols = self.return_all_event_cols,
+                    return_nsrrid = self.return_nsrrid,
+                )
+            self._train_dataset = train_set
+            persistent_workers = True
+        
+        if self.n_train_samples is not None and self.n_train_samples > 0:
+            n_total = len(train_set)
+            rng = np.random.default_rng(seed=self.random_seed)
+            
+            if hasattr(train_set, 'event_cols') and train_set.event_cols and hasattr(train_set, 'all_epoch_df'):
+                label_col = train_set.event_cols[0]
+                if label_col in train_set.all_epoch_df.columns:
+                    labels = train_set.all_epoch_df[label_col].values
+                    num_classes = getattr(train_set, 'num_classes', None)
+                    
+                    if num_classes is not None:
+                        all_indices = []
+                        for c in range(num_classes):
+                            class_indices = np.where(labels == c)[0]
+                            n_per_class = min(self.n_train_samples, len(class_indices))
+                            if n_per_class > 0:
+                                sampled = rng.choice(class_indices, size=n_per_class, replace=False)
+                                all_indices.extend(sampled.tolist())
+                                print(f"[Few-shot] Class {c}: sampled {n_per_class}/{len(class_indices)} samples")
+                        
+                        indices = all_indices
+                        train_set = Subset(train_set, indices)
+                        print(f"[Few-shot] Total: {len(indices)}/{n_total} samples ({self.n_train_samples}-shot per class)")
+                    else:
+                        n_keep = min(self.n_train_samples, n_total)
+                        indices = rng.choice(n_total, size=n_keep, replace=False).tolist()
+                        train_set = Subset(train_set, indices)
+                        print(f"[Few-shot] Using {n_keep}/{n_total} training samples (random, n_train_samples={self.n_train_samples})")
+                else:
+                    n_keep = min(self.n_train_samples, n_total)
+                    indices = rng.choice(n_total, size=n_keep, replace=False).tolist()
+                    train_set = Subset(train_set, indices)
+                    print(f"[Few-shot] Using {n_keep}/{n_total} training samples (random, n_train_samples={self.n_train_samples})")
+            else:
+                n_keep = min(self.n_train_samples, n_total)
+                indices = rng.choice(n_total, size=n_keep, replace=False).tolist()
+                train_set = Subset(train_set, indices)
+                print(f"[Few-shot] Using {n_keep}/{n_total} training samples (random, n_train_samples={self.n_train_samples})")
+        
+        return DataLoader(
+            train_set,
+            batch_size     = self.hparams.batch_size,
+            shuffle        = True,
+            num_workers    = self.hparams.num_workers,
+            pin_memory     = self.pin_memory,
+            persistent_workers = persistent_workers,
+            drop_last      = True,
+        )
+    
+    def get_class_distribution(self) -> Optional[torch.Tensor]:
+        """
+        Get class distribution from training dataset.
+        Returns [num_classes] tensor of class counts, or None if not available.
+        """
+        if hasattr(self, '_train_dataset'):
+            counts = self._train_dataset.get_class_counts()
+            if counts is not None:
+                return torch.from_numpy(counts).float()
+        return None
+
+    def val_dataloader(self):
+        if self.hparams.val_dataset_list:       
+            if self.is_pretrain == 1:
+                val_sets = [
+                        SleepEpochDataset(
+                            csv_dir       = self.csv_dir,
+                            split         = "pretrain-val",
+                            data_pct      = self.data_pct,
+                            patient_cols   = self.patient_cols,
+                            event_cols   = self.event_cols,
+                            train_edf_cols= self.hparams.train_edf_cols,
+                            transform     = None,        
+                            sample_rate   = self.hparams.sample_rate,
+                            window_size   = self.hparams.window_size,
+                            cache_size    = self.hparams.cache_size,
+                            downstream_dataset_name  = ds_name,
+                            data_source   = self.data_source,
+                            include_datasets = self.include_datasets,
+                        )
+                        for ds_name in self.hparams.val_dataset_list
+                    ]
+                persistent_workers = self.persistent_workers
+        else:
+            if self.is_pretrain == 1:
+                val_sets = [
+                    SleepEpochDataset(
+                        csv_dir       = self.csv_dir,
+                        split         = "pretrain-val",
+                        data_pct      = self.data_pct,
+                        patient_cols   = self.patient_cols,
+                        event_cols   = self.event_cols,
+                        train_edf_cols= self.hparams.train_edf_cols,
+                        transform     = None,
+                        sample_rate   = self.hparams.sample_rate,
+                        window_size   = self.hparams.window_size,
+                        cache_size    = self.hparams.cache_size,
+                        data_source   = self.data_source,
+                        include_datasets = self.include_datasets,
+                    )
+                    ]
+                persistent_workers = self.persistent_workers
+            else:
+                val_sets = [
+                    SleepEpochDataset(
+                        csv_dir       = self.csv_dir,
+                        split         = "val",
+                        data_pct      = self.data_pct,
+                        patient_cols   = self.patient_cols,
+                        event_cols   = self.event_cols,
+                        train_edf_cols= self.hparams.train_edf_cols,
+                        transform     = None,
+                        sample_rate   = self.hparams.sample_rate,
+                        window_size   = self.hparams.window_size,
+                        cache_size    = self.hparams.cache_size,
+                        downstream_dataset_name  = self.downstream_dataset_name,
+                        data_source   = self.data_source,
+                        include_datasets = self.include_datasets,
+                        regression_targets = self.regression_targets,
+                        regression_filter_config = self.regression_filter_config,
+                    )
+                    ]
+                persistent_workers = True
+        
+        if self.val_data_pct is not None and 0 < self.val_data_pct < 1.0:
+            subsampled_val_sets = []
+            for ds in val_sets:
+                n_total = len(ds)
+                n_keep = max(1, int(n_total * self.val_data_pct))
+                rng = np.random.default_rng(seed=self.random_seed)
+                indices = rng.choice(n_total, size=n_keep, replace=False).tolist()
+                subsampled_val_sets.append(Subset(ds, indices))
+                print(f"[Val subsample] Using {n_keep}/{n_total} val samples ({self.val_data_pct*100:.1f}%)")
+            val_sets = subsampled_val_sets
+        
+        val_bs = self.val_batch_size if self.val_batch_size is not None else self.hparams.batch_size
+        return [
+            DataLoader(
+                ds,
+                batch_size     = val_bs,
+                shuffle        = False,
+                num_workers    = self.hparams.num_workers,
+                pin_memory     = self.pin_memory,
+                persistent_workers = persistent_workers,
+                drop_last      = True,
+            )
+            for ds in val_sets
+        ]
+
+    def test_dataloader(self):
+        if self.is_pretrain == 1:
+            test_set = SleepEpochDataset(
+                csv_dir       = self.csv_dir,
+                split         = "pretrain-test",
+                patient_cols   = self.patient_cols,
+                event_cols   = self.event_cols,
+                train_edf_cols= self.hparams.train_edf_cols,
+                transform     = None,
+                sample_rate   = self.hparams.sample_rate,
+                window_size   = self.hparams.window_size,
+                cache_size    = self.hparams.cache_size,
+                data_source   = self.data_source,
+                include_datasets = self.include_datasets,
+            )
+            persistent_workers = self.persistent_workers
+        else:
+            test_set = SleepEpochDataset(
+                    csv_dir       = self.csv_dir,
+                    split         = "test",
+                    patient_cols   = self.patient_cols,
+                    event_cols   = self.event_cols,
+                    train_edf_cols= self.hparams.train_edf_cols,
+                    transform     = None,
+                    sample_rate   = self.hparams.sample_rate,
+                    window_size   = self.hparams.window_size,
+                    cache_size    = self.hparams.cache_size,
+                    downstream_dataset_name  = self.downstream_dataset_name,
+                    data_source   = self.data_source,
+                    include_datasets = self.include_datasets,
+                    regression_targets = self.regression_targets,
+                    regression_filter_config = self.regression_filter_config,
+                )
+            persistent_workers = True
+        test_bs = self.val_batch_size if self.val_batch_size is not None else self.hparams.batch_size
+        return DataLoader(
+            test_set,
+            batch_size     = test_bs,
+            shuffle        = False,
+            num_workers    = self.hparams.num_workers,
+            pin_memory     = self.pin_memory,
+            drop_last      = True,
+            persistent_workers = persistent_workers,
+        )
+

osf/datasets/pretrain_dataset.py

@@ -0,0 +1,381 @@
+# Sleep Epoch Dataset for pretraining and downstream tasks
+
+import os
+import numpy as np
+import pandas as pd
+import torch
+from pathlib import Path
+from contextlib import suppress
+from typing import Sequence, Optional, Dict, Union, List
+from torch.utils.data import Dataset
+from train_config import NEED_NORM_COL
+
+
+def to_pm1(s: pd.Series) -> pd.Series:
+    s = pd.to_numeric(s, errors="coerce")
+    vmin, vmax = s.min(skipna=True), s.max(skipna=True)
+    if pd.isna(vmin) or pd.isna(vmax) or vmax <= vmin:
+        return pd.Series(0.0, index=s.index)
+    return (2 * (s - vmin) / (vmax - vmin) - 1).fillna(0.0)
+
+
+class SleepEpochDataset(Dataset):
+    def __init__(
+        self,
+        csv_dir='/path/to/your/postprocessed/data',
+        split: str = "train",
+        *,
+        data_pct=1,
+        patient_cols: Optional[Union[str, Sequence[str]]] = None,
+        event_cols: Optional[Union[str, Sequence[str]]] = None,
+        train_edf_cols=None,
+        test_size: float = 0.15,
+        val_size: float = 0.15,
+        random_state: int = 1337,
+        sample_rate: int = 128,
+        window_size: int = 300,
+        epoch_length: int = 30,
+        cache_size: int = 8,
+        transform=None,
+        downstream_dataset_name=None,
+        data_source: str = "auto",
+        include_datasets: Optional[List[str]] = None,
+        regression_targets: Optional[List[str]] = None,
+        regression_filter_config: Optional[Dict] = None,
+        return_all_event_cols: bool = False,
+        return_nsrrid: bool = False,
+    ):
+        assert split in {"pretrain", "pretrain-val", "pretrain-test", "train", "val", "test"}
+        assert data_source in {"auto", "pretrain", "downstream", "both"}
+
+        self.transform = transform
+        self.sample_rate = sample_rate
+        self.window_size = window_size
+        self.epoch_length = epoch_length
+        self.patient_cols = [patient_cols] if isinstance(patient_cols, str) else patient_cols
+        self.event_cols = [event_cols] if isinstance(event_cols, str) else event_cols
+        self.train_edf_cols = train_edf_cols
+        self.split = split
+        self.data_pct = float(data_pct)
+        self.data_source = data_source
+        self.regression_targets = regression_targets
+        self.regression_filter_config = regression_filter_config
+        self.return_all_event_cols = return_all_event_cols
+        self.return_nsrrid = return_nsrrid
+
+        patient_df, epoch_df = self._load_csvs(
+            csv_dir, split, data_source, include_datasets, self.event_cols,
+            regression_targets=self.regression_targets,
+            regression_filter_config=self.regression_filter_config,
+            return_all_event_cols=self.return_all_event_cols,
+        )
+
+        if downstream_dataset_name and include_datasets is None:
+            if downstream_dataset_name != "all":
+                mask = epoch_df['dataset_name'].astype(str).str.lower().str.startswith(downstream_dataset_name)
+                epoch_df = epoch_df.loc[mask].copy()
+                ids = epoch_df["nsrrid"].astype(str).unique()
+                patient_df = patient_df[patient_df["nsrrid"].astype(str).isin(ids)].copy()
+
+        # Determine num_classes
+        if self.event_cols:
+            if self.event_cols[0] in ['Hypopnea', 'Arousal', 'Oxygen Desaturation']:
+                self.num_classes = 2
+            elif self.event_cols[0] == 'Stage':
+                self.num_classes = 4
+                mapping = {0: 0, 1: 1, 2: 1, 3: 2, 4: 3}
+                epoch_df['Stage'] = epoch_df['Stage'].replace(mapping)
+            else:
+                self.num_classes = 2
+        else:
+            self.num_classes = 2
+
+        # Drop Stage == -1
+        if self.event_cols and ('Stage' in self.event_cols) and ('Stage' in epoch_df.columns):
+            epoch_df = epoch_df.loc[epoch_df['Stage'] != -1].copy()
+
+        # Build tables
+        if split in ("pretrain", "pretrain-val"):
+            sort_cols = [c for c in ['nsrrid', 'seg_id', 'epoch_id'] if c in epoch_df.columns]
+            self.all_epoch_df = epoch_df.sort_values(sort_cols).reset_index(drop=True)
+
+            idx_keep_cols = [c for c in ['nsrrid', 'seg_id', 'path_head'] if c in self.all_epoch_df.columns]
+            if self.regression_targets:
+                for t in self.regression_targets:
+                    col = f"{t}_mean"
+                    if col in self.all_epoch_df.columns:
+                        idx_keep_cols.append(col)
+            self.epoch_df = (
+                self.all_epoch_df[idx_keep_cols]
+                .drop_duplicates(['nsrrid', 'seg_id'], keep='first')
+                .reset_index(drop=True)
+            )
+        else:
+            expected_len = self.window_size // self.epoch_length
+            grp = epoch_df.groupby(['nsrrid', 'seg_id']).size().rename('n').reset_index()
+            valid_keys = grp.loc[grp['n'] == expected_len, ['nsrrid', 'seg_id']]
+            epoch_df_valid = epoch_df.merge(valid_keys, on=['nsrrid', 'seg_id'], how='inner')
+
+            sort_cols = [c for c in ['nsrrid', 'seg_id', 'epoch_id'] if c in epoch_df_valid.columns]
+            self.all_epoch_df = epoch_df_valid.sort_values(sort_cols).reset_index(drop=True)
+
+            idx_keep_cols = [c for c in ['nsrrid', 'seg_id', 'path_head'] if c in self.all_epoch_df.columns]
+            if self.regression_targets:
+                for t in self.regression_targets:
+                    col = f"{t}_mean"
+                    if col in self.all_epoch_df.columns:
+                        idx_keep_cols.append(col)
+            self.epoch_df = (
+                self.all_epoch_df[idx_keep_cols]
+                .drop_duplicates(['nsrrid', 'seg_id'], keep='first')
+                .reset_index(drop=True)
+            )
+
+        # Patient-level sampling
+        if not (0 < self.data_pct <= 1.0):
+            raise ValueError(f"data_pct must be in (0,1], got {self.data_pct}")
+
+        if self.data_pct < 1.0:
+            eligible_patients = pd.Index(self.epoch_df['nsrrid'].unique())
+            n_keep = max(1, int(len(eligible_patients) * self.data_pct))
+            sampled_nsrrids = pd.Series(eligible_patients).sample(n=n_keep, random_state=random_state).to_list()
+            self.epoch_df = self.epoch_df.loc[self.epoch_df['nsrrid'].isin(sampled_nsrrids)].reset_index(drop=True)
+            self.all_epoch_df = self.all_epoch_df.loc[self.all_epoch_df['nsrrid'].isin(sampled_nsrrids)].reset_index(drop=True)
+            patient_df = patient_df.loc[patient_df['nsrrid'].isin(sampled_nsrrids)].copy()
+
+        self.patient_df = patient_df.set_index("nsrrid")
+
+        # Build segment indices
+        self._seg_indices = None
+        if hasattr(self, "all_epoch_df") and {'nsrrid', 'seg_id'}.issubset(self.all_epoch_df.columns):
+            grp_indices = self.all_epoch_df.groupby(['nsrrid', 'seg_id'], sort=False).indices
+            self._seg_indices = {}
+            has_epoch_id = 'epoch_id' in self.all_epoch_df.columns
+            epoch_id_values = self.all_epoch_df['epoch_id'].to_numpy() if has_epoch_id else None
+            for key, idx_list in grp_indices.items():
+                idx_arr = np.fromiter(idx_list, dtype=np.int64)
+                if has_epoch_id:
+                    order = np.argsort(epoch_id_values[idx_arr])
+                    idx_arr = idx_arr[order]
+                self._seg_indices[key] = idx_arr
+
+        # Compute class distribution
+        self._class_counts = None
+        if self.event_cols and self.event_cols[0] in self.all_epoch_df.columns:
+            label_col = self.event_cols[0]
+            value_counts = self.all_epoch_df[label_col].value_counts().sort_index()
+            class_counts = np.zeros(self.num_classes, dtype=np.int64)
+            for cls_idx, count in value_counts.items():
+                if 0 <= int(cls_idx) < self.num_classes:
+                    class_counts[int(cls_idx)] = int(count)
+            self._class_counts = class_counts
+
+    def _load_csvs(self, csv_dir, split, data_source, include_datasets, event_cols,
+                   regression_targets=None, regression_filter_config=None, return_all_event_cols=False):
+        split_suffix_map = {
+            "pretrain": "train", "pretrain-val": "valid", "pretrain-test": "test",
+            "train": "train", "val": "valid", "test": "test"
+        }
+        split_suffix = split_suffix_map[split]
+
+        if data_source == "auto":
+            sources = ["pretrain"] if split.startswith("pretrain") else ["downstream"]
+        elif data_source == "both":
+            sources = ["pretrain", "downstream"]
+        else:
+            sources = [data_source]
+
+        patient_dfs = []
+        epoch_dfs = []
+        csv_prefix = "epoch_regression" if regression_targets else "epoch"
+
+        for source in sources:
+            patient_csv = f"{csv_dir}/patient_{source}_{split_suffix}.csv"
+            epoch_csv = f"{csv_dir}/{csv_prefix}_{source}_{split_suffix}.csv"
+
+            if Path(patient_csv).is_file() and Path(epoch_csv).is_file():
+                patient_dfs.append(pd.read_csv(patient_csv))
+                epoch_dfs.append(pd.read_csv(epoch_csv))
+
+        patient_df = pd.concat(patient_dfs, ignore_index=True).drop_duplicates(subset=['nsrrid'])
+        epoch_df = pd.concat(epoch_dfs, ignore_index=True)
+
+        base_cols = ['nsrrid', 'seg_id', 'dataset_name', 'epoch_id', 'path_head']
+        if event_cols:
+            if return_all_event_cols:
+                for col in event_cols:
+                    if col and col not in base_cols:
+                        base_cols.append(col)
+            elif event_cols[0]:
+                base_cols.append(event_cols[0])
+
+        if regression_targets:
+            for t in regression_targets:
+                col_name = f"{t}_mean"
+                if col_name in epoch_df.columns:
+                    base_cols.append(col_name)
+
+        keep_cols = [c for c in base_cols if c in epoch_df.columns]
+        epoch_df = epoch_df[keep_cols].copy()
+
+        if regression_targets:
+            label_cols = [f"{t}_mean" for t in regression_targets]
+            existing = [c for c in label_cols if c in epoch_df.columns]
+            if existing:
+                epoch_df = epoch_df.dropna(subset=existing).reset_index(drop=True)
+
+        if regression_filter_config:
+            for col_name, filter_rules in regression_filter_config.items():
+                if col_name in epoch_df.columns:
+                    mask = pd.Series([True] * len(epoch_df))
+                    if "min" in filter_rules:
+                        mask = mask & (epoch_df[col_name] >= filter_rules["min"])
+                    if "max" in filter_rules:
+                        mask = mask & (epoch_df[col_name] <= filter_rules["max"])
+                    epoch_df = epoch_df[mask].reset_index(drop=True)
+
+        if include_datasets is not None and 'dataset_name' in epoch_df.columns:
+            include_lower = [d.lower() for d in include_datasets]
+            mask = epoch_df['dataset_name'].astype(str).str.lower().isin(include_lower)
+            epoch_df = epoch_df[mask].copy()
+            patient_df = patient_df[patient_df['nsrrid'].isin(epoch_df['nsrrid'].unique())].copy()
+
+        return patient_df, epoch_df
+
+    def __len__(self) -> int:
+        return len(self.epoch_df)
+
+    def get_class_counts(self) -> Optional[np.ndarray]:
+        return self._class_counts
+
+    def _resample_df(self, df: pd.DataFrame, target_hz: int) -> pd.DataFrame:
+        if not np.issubdtype(df.index.dtype, np.number):
+            t = np.arange(len(df)) / float(target_hz)
+            df = df.copy()
+            df.index = t
+
+        t0 = float(df.index.min())
+        t1 = float(df.index.max())
+        t_target = np.arange(t0, t0 + self.window_size, 1.0 / target_hz)
+        if t_target[-1] > t1:
+            t_target = t_target[t_target <= t1 + 1e-9]
+        out = df.reindex(t_target).interpolate(method="linear", limit_direction="both")
+        return out.fillna(0.0)
+
+    def __getitem__(self, idx: int):
+        row = self.epoch_df.iloc[idx]
+        nsrrid = row["nsrrid"]
+        seg_id = int(row["seg_id"])
+        cols = list(self.train_edf_cols) if self.train_edf_cols is not None else None
+
+        if self.split == "pretrain":
+            df_epoch = self._load_epoch_all_df(row["path_head"], seg_id, columns=cols)
+            df_epoch = self._resample_df(df_epoch, self.sample_rate)
+
+            if cols is not None:
+                for ch in cols:
+                    if ch not in df_epoch.columns:
+                        df_epoch[ch] = 0.0
+                    elif ch in NEED_NORM_COL:
+                        df_epoch[ch] = to_pm1(df_epoch[ch])
+                df_epoch = df_epoch[cols]
+
+            samples_per_epoch = int(self.window_size * self.sample_rate)
+            if len(df_epoch) < samples_per_epoch:
+                pad = samples_per_epoch - len(df_epoch)
+                tail = pd.DataFrame({c: 0.0 for c in df_epoch.columns},
+                                    index=df_epoch.index[-1] + (np.arange(1, pad + 1) / self.sample_rate))
+                df_epoch = pd.concat([df_epoch, tail], axis=0)
+            elif len(df_epoch) > samples_per_epoch:
+                df_epoch = df_epoch.iloc[:samples_per_epoch]
+
+            x = torch.tensor(df_epoch.to_numpy(copy=False), dtype=torch.float32).t().contiguous()
+            x = torch.clamp(x, min=-6, max=6)
+
+            output = {"psg": x}
+            if self.return_nsrrid:
+                output["nsrrid"] = nsrrid
+                output["seg_id"] = seg_id
+
+            if self.patient_cols:
+                y = torch.tensor(self.patient_df.loc[nsrrid, self.patient_cols].values.astype(float), dtype=torch.float32)
+                output["label"] = y.long() if not self.return_nsrrid else y
+            elif self.event_cols:
+                if self.return_all_event_cols:
+                    available_cols = [c for c in self.event_cols if c in row.index]
+                    y = torch.tensor([row[c] for c in available_cols], dtype=torch.float32)
+                else:
+                    y = torch.tensor([row[self.event_cols[0]]], dtype=torch.float32)
+                output["label"] = y
+
+            return output
+        else:
+            # Downstream split
+            if self._seg_indices is None:
+                seg_df = self.all_epoch_df[
+                    (self.all_epoch_df['nsrrid'] == nsrrid) & (self.all_epoch_df['seg_id'] == seg_id)
+                ].sort_values('epoch_id')
+            else:
+                idx_arr = self._seg_indices.get((nsrrid, seg_id))
+                seg_df = self.all_epoch_df.iloc[idx_arr] if idx_arr is not None else \
+                    self.all_epoch_df[(self.all_epoch_df['nsrrid'] == nsrrid) & (self.all_epoch_df['seg_id'] == seg_id)].sort_values('epoch_id')
+
+            df_epoch = self._load_epoch_all_df(row["path_head"], seg_id, columns=cols)
+            df_epoch = self._resample_df(df_epoch, self.sample_rate)
+
+            if cols is not None:
+                for ch in cols:
+                    if ch not in df_epoch.columns:
+                        df_epoch[ch] = 0.0
+                    elif ch in NEED_NORM_COL:
+                        df_epoch[ch] = to_pm1(df_epoch[ch])
+                df_epoch = df_epoch[cols]
+
+            samples_per_epoch = int(self.window_size * self.sample_rate)
+            if len(df_epoch) < samples_per_epoch:
+                pad = samples_per_epoch - len(df_epoch)
+                tail = pd.DataFrame({c: 0.0 for c in df_epoch.columns},
+                                    index=df_epoch.index[-1] + (np.arange(1, pad + 1) / self.sample_rate))
+                df_epoch = pd.concat([df_epoch, tail], axis=0)
+            elif len(df_epoch) > samples_per_epoch:
+                df_epoch = df_epoch.iloc[:samples_per_epoch]
+
+            x = torch.tensor(df_epoch.to_numpy(copy=False), dtype=torch.float32).t().contiguous()
+            x = torch.clamp(x, min=-6, max=6)
+
+            output = {"psg": x}
+            if self.return_nsrrid:
+                output["nsrrid"] = nsrrid
+                output["seg_id"] = seg_id
+
+            if self.patient_cols:
+                y = torch.tensor(self.patient_df.loc[nsrrid, self.patient_cols].values.astype(float), dtype=torch.float32)
+                y = y.repeat(self.window_size // self.epoch_length)
+                output["label"] = y
+            elif self.event_cols:
+                if self.return_all_event_cols:
+                    available_cols = [c for c in self.event_cols if c in seg_df.columns]
+                    y = torch.tensor(seg_df[available_cols].values.astype(float), dtype=torch.float32).squeeze(0)
+                else:
+                    y = torch.tensor(seg_df[self.event_cols].values.astype(float), dtype=torch.float32).squeeze(1)
+                output["label"] = y
+            elif self.regression_targets:
+                label_cols = [f"{t}_mean" for t in self.regression_targets]
+                y = torch.tensor([row[c] for c in label_cols], dtype=torch.float32)
+                output["label"] = y
+
+            return output
+
+    def _build_epoch_all_path(self, path_head: str, epoch_id: int) -> Path:
+        return Path(f"{path_head}/epoch-{epoch_id:05d}_all.parquet")
+
+    def _load_epoch_all_df(self, path_head: str, epoch_id: int, columns=None) -> pd.DataFrame:
+        fp = self._build_epoch_all_path(path_head, epoch_id)
+        if not fp.is_file():
+            raise FileNotFoundError(f"Parquet missing: {fp}")
+        df = pd.read_parquet(fp)
+        for c in df.columns:
+            if not np.issubdtype(df[c].dtype, np.floating):
+                with suppress(Exception):
+                    df[c] = df[c].astype(np.float32)
+        return df

osf/datasets/simclr_aug_registry.py

@@ -0,0 +1,258 @@
+"""
+Two-view augmentation registry for SSL pretraining (SimCLR, DINO).
+Provides multi-view generation pipelines for contrastive and self-distillation methods.
+"""
+from __future__ import annotations
+from typing import Callable, Dict
+import torch
+from osf.datasets import augmentations as A
+
+
+def _two_view(pipe1: Callable, pipe2: Callable | None = None) -> Callable:
+    """Wrap one/two single-view pipelines into a two-view augmentation maker."""
+    if pipe2 is None:
+        pipe2 = pipe1
+    def make(x: torch.Tensor):
+        return pipe1(x), pipe2(x)
+    return make
+
+
+SIMCLR_AUG_REGISTRY: Dict[str, Callable] = {
+    "none": _two_view(lambda x: x),
+    
+    "channel_dropout": _two_view(lambda x: A.channel_dropout(x, drop_prob=0.2, min_keep=1)),
+    "channel_dropout_light": _two_view(lambda x: A.channel_dropout(x, drop_prob=0.25, min_keep=1)),
+    "channel_dropout_aligned": _two_view(lambda x: A.channel_dropout(x, drop_prob=0.5, min_keep=1)),
+}
+
+SIMCLR_AUG_FACTORIES: Dict[str, Callable[..., Callable]] = {}
+
+
+def build_simclr_augmentor(name: str, **kwargs) -> Callable:
+    key = (name or "none").lower()
+    if key in SIMCLR_AUG_REGISTRY:
+        return SIMCLR_AUG_REGISTRY[key]
+    if key in SIMCLR_AUG_FACTORIES:
+        return SIMCLR_AUG_FACTORIES[key](**kwargs)
+    raise ValueError(
+        f"Unknown simclr_augmentation '{name}'. "
+        f"Available presets: {list(SIMCLR_AUG_REGISTRY.keys())} | "
+        f"factories: {list(SIMCLR_AUG_FACTORIES.keys())}"
+    )
+
+
+def _per_channel_span_mask_factory(
+    ratio: tuple[float, float] = (0.10, 0.30),
+    n_spans: int = 1,
+    fill: str | torch.Tensor = "zero",
+    noise_scale: float = 0.05,
+    same_mask_for_batch: bool = False,
+):
+    assert 0.0 <= ratio[0] <= ratio[1] <= 1.0
+
+    def _single_view(x: torch.Tensor) -> torch.Tensor:
+        B, C, T = x.shape
+        device, dtype = x.device, x.dtype
+
+        min_len = max(1, int(round(ratio[0] * T)))
+        max_len = max(min_len, int(round(ratio[1] * T)))
+        arange_T = torch.arange(T, device=device)
+        mask = torch.zeros((B, C, T), device=device, dtype=torch.bool)
+        shape_bc = (1, C) if same_mask_for_batch else (B, C)
+
+        for _ in range(max(1, int(n_spans))):
+            if max_len == min_len:
+                lengths = torch.full(shape_bc, max_len, device=device, dtype=torch.long)
+            else:
+                lengths = torch.randint(min_len, max_len + 1, shape_bc, device=device)
+            max_start = (T - lengths).clamp_min(0)
+            if (max_start > 0).any():
+                rnd = torch.rand_like(max_start, dtype=torch.float32)
+                starts = torch.floor(rnd * (max_start.to(torch.float32) + 1)).to(torch.long)
+            else:
+                starts = torch.zeros_like(max_start)
+            if same_mask_for_batch and B > 1:
+                starts = starts.expand(B, C)
+                lengths = lengths.expand(B, C)
+            span_mask = (arange_T.view(1, 1, T) >= starts.unsqueeze(-1)) & \
+                        (arange_T.view(1, 1, T) < (starts + lengths).unsqueeze(-1))
+            mask |= span_mask
+
+        y = x.clone()
+        if isinstance(fill, torch.Tensor):
+            fill_t = fill.to(device=device, dtype=dtype)
+            if fill_t.dim() == 0:
+                fill_t = fill_t.view(1, 1, 1)
+            if fill_t.shape[-1] == 1 and fill_t.dim() == 3 and fill_t.shape[0] in (1, B):
+                fill_t = fill_t if fill_t.shape[0] == B else fill_t.expand(B, -1, -1)
+            elif fill_t.dim() == 3 and fill_t.shape == (B, C, T):
+                pass
+            elif fill_t.dim() == 3 and fill_t.shape == (1, C, 1):
+                fill_t = fill_t.expand(B, -1, T)
+            y[mask] = fill_t[mask.expand_as(fill_t)]
+        elif fill == "zero":
+            y[mask] = 0.0
+        elif fill == "mean":
+            m = x.mean(dim=-1, keepdim=True)
+            y = torch.where(mask, m.expand_as(x), y)
+        elif fill == "noise":
+            m = x.mean(dim=-1, keepdim=True)
+            s = x.std(dim=-1, keepdim=True, unbiased=False).clamp_min(1e-8)
+            noise = torch.randn_like(x) * (s * noise_scale) + m
+            y = torch.where(mask, noise, y)
+        else:
+            raise ValueError(f"Unknown fill mode: {fill!r}")
+        return y
+
+    return _two_view(_single_view)
+
+
+SIMCLR_AUG_FACTORIES["pc_span_mask"] = _per_channel_span_mask_factory
+
+SIMCLR_AUG_REGISTRY.update({
+    "pc_span_mask_light": _per_channel_span_mask_factory(
+        ratio=(0.1, 0.3), n_spans=1, fill="zero", noise_scale=0.05, same_mask_for_batch=False
+    ),
+    "pc_span_mask_heavy": _per_channel_span_mask_factory(
+        ratio=(0.20, 0.6), n_spans=2, fill="zero", noise_scale=0.05, same_mask_for_batch=False
+    ),
+    "pc_span_mask_aligned": _per_channel_span_mask_factory(
+        ratio=(0.3, 0.6), n_spans=1, fill="zero", noise_scale=0, same_mask_for_batch=False
+    ),
+})
+
+
+def _channel_then_pcspan_factory(
+    drop_prob: float = 0.3,
+    min_keep: int = 1,
+    ratio: tuple[float, float] = (0.10, 0.30),
+    n_spans: int = 1,
+    fill: str = "zero",
+    noise_scale: float = 0.05,
+    same_mask_for_batch: bool = False,
+):
+    def single_view(x: torch.Tensor) -> torch.Tensor:
+        y = A.channel_dropout(x, drop_prob=drop_prob, min_keep=min_keep)
+        B, C, T = y.shape
+        device = y.device
+
+        min_len = max(1, int(round(ratio[0] * T)))
+        max_len = max(min_len, int(round(ratio[1] * T)))
+        arange_T = torch.arange(T, device=device)
+        mask = torch.zeros((B, C, T), device=device, dtype=torch.bool)
+        shape_bc = (1, C) if same_mask_for_batch else (B, C)
+
+        for _ in range(max(1, int(n_spans))):
+            lengths = torch.full(shape_bc, max_len, device=device, dtype=torch.long) \
+                      if max_len == min_len else torch.randint(min_len, max_len + 1, shape_bc, device=device)
+            max_start = (T - lengths).clamp_min(0)
+            if (max_start > 0).any():
+                rnd = torch.rand_like(max_start, dtype=torch.float32)
+                starts = torch.floor(rnd * (max_start.to(torch.float32) + 1)).to(torch.long)
+            else:
+                starts = torch.zeros_like(max_start)
+            if same_mask_for_batch and B > 1:
+                starts = starts.expand(B, C)
+                lengths = lengths.expand(B, C)
+            span_mask = (arange_T.view(1, 1, T) >= starts.unsqueeze(-1)) & \
+                        (arange_T.view(1, 1, T) < (starts + lengths).unsqueeze(-1))
+            mask |= span_mask
+
+        out = y.clone()
+        if fill == "zero":
+            out[mask] = 0.0
+        elif fill == "mean":
+            m = y.mean(dim=-1, keepdim=True)
+            out = torch.where(mask, m.expand_as(y), out)
+        elif fill == "noise":
+            m = y.mean(dim=-1, keepdim=True)
+            s = y.std(dim=-1, keepdim=True, unbiased=False).clamp_min(1e-8)
+            noise = torch.randn_like(y) * (s * noise_scale) + m
+            out = torch.where(mask, noise, out)
+        else:
+            raise ValueError(f"Unknown fill: {fill!r}")
+        return out
+
+    return _two_view(single_view)
+
+
+SIMCLR_AUG_FACTORIES["chan_then_pcspan"] = _channel_then_pcspan_factory
+
+
+def _crop_then_chan_pcspan_factory(
+    crop_ratio: tuple[float, float] = (0.25, 0.75),
+    align_to: int = 40,
+    drop_prob: float = 0.5,
+    min_keep: int = 1,
+    span_ratio: tuple[float, float] = (0.3, 0.6),
+    n_spans: int = 1,
+    fill: str = "zero",
+    noise_scale: float = 0.0,
+    same_mask_for_batch: bool = False,
+):
+    def single_view(x: torch.Tensor) -> torch.Tensor:
+        y = A.random_time_crop(x, ratio=crop_ratio, resize_back=True, align_to=align_to)
+        y = A.channel_dropout(y, drop_prob=drop_prob, min_keep=min_keep)
+
+        B, C, T = y.shape
+        device = y.device
+        min_len = max(1, int(round(span_ratio[0] * T)))
+        max_len = max(min_len, int(round(span_ratio[1] * T)))
+        arange_T = torch.arange(T, device=device)
+        mask = torch.zeros((B, C, T), device=device, dtype=torch.bool)
+        shape_bc = (1, C) if same_mask_for_batch else (B, C)
+
+        for _ in range(max(1, int(n_spans))):
+            lengths = torch.full(shape_bc, max_len, device=device, dtype=torch.long) \
+                      if max_len == min_len else torch.randint(min_len, max_len + 1, shape_bc, device=device)
+            max_start = (T - lengths).clamp_min(0)
+            if (max_start > 0).any():
+                rnd = torch.rand_like(max_start, dtype=torch.float32)
+                starts = torch.floor(rnd * (max_start.to(torch.float32) + 1)).to(torch.long)
+            else:
+                starts = torch.zeros_like(max_start)
+            if same_mask_for_batch and B > 1:
+                starts = starts.expand(B, C)
+                lengths = lengths.expand(B, C)
+            span_mask = (arange_T.view(1, 1, T) >= starts.unsqueeze(-1)) & \
+                        (arange_T.view(1, 1, T) < (starts + lengths).unsqueeze(-1))
+            mask |= span_mask
+
+        out = y.clone()
+        if fill == "zero":
+            out[mask] = 0.0
+        elif fill == "mean":
+            m = y.mean(dim=-1, keepdim=True)
+            out = torch.where(mask, m.expand_as(y), out)
+        elif fill == "noise":
+            m = y.mean(dim=-1, keepdim=True)
+            s = y.std(dim=-1, keepdim=True, unbiased=False).clamp_min(1e-8)
+            noise = torch.randn_like(y) * (s * noise_scale) + m
+            out = torch.where(mask, noise, out)
+        else:
+            raise ValueError(f"Unknown fill: {fill!r}")
+        return out
+
+    return _two_view(single_view)
+
+
+SIMCLR_AUG_FACTORIES["crop_then_chan_pcspan"] = _crop_then_chan_pcspan_factory
+
+SIMCLR_AUG_REGISTRY.update({
+    "chan_then_pcspan": _channel_then_pcspan_factory(
+        drop_prob=0.5, min_keep=1, ratio=(0.3, 0.6), n_spans=1, fill="zero",
+        noise_scale=0, same_mask_for_batch=False
+    ),
+    "chan_then_pcspan_light": _channel_then_pcspan_factory(
+        drop_prob=0.25, min_keep=1, ratio=(0.3, 0.6), n_spans=1, fill="zero",
+        noise_scale=0, same_mask_for_batch=False
+    ),
+    "crop_then_chan_pcspan": _crop_then_chan_pcspan_factory(
+        crop_ratio=(0.25, 0.75), align_to=40, drop_prob=0.5, min_keep=1,
+        span_ratio=(0.3, 0.6), n_spans=1, fill="zero", noise_scale=0, same_mask_for_batch=False
+    ),
+    "crop_then_chan_pcspan_light": _crop_then_chan_pcspan_factory(
+        crop_ratio=(0.25, 0.75), align_to=40, drop_prob=0.25, min_keep=1,
+        span_ratio=(0.3, 0.6), n_spans=1, fill="zero", noise_scale=0, same_mask_for_batch=False
+    ),
+})

osf/models/balanced_losses.py

@@ -0,0 +1,89 @@
+"""
+Balanced/imbalanced learning losses.
+Reference: https://github.com/YyzHarry/SubpopBench
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+
+class FocalLoss(nn.Module):
+    """
+    Focal Loss: FL(p_t) = -α_t * (1 - p_t)^γ * log(p_t)
+    Paper: https://arxiv.org/abs/1708.02002
+    
+    Args:
+        alpha: Weighting factor (float or [num_classes] tensor)
+        gamma: Focusing parameter (higher = more focus on hard examples)
+        reduction: 'mean' or 'none'
+    """
+    def __init__(self, alpha: Optional[float | torch.Tensor] = None, gamma: float = 2.0, reduction: str = "mean"):
+        super().__init__()
+        self.gamma = gamma
+        self.reduction = reduction
+        
+        if isinstance(alpha, (float, int)):
+            self.register_buffer("alpha", torch.tensor([alpha], dtype=torch.float32))
+        elif isinstance(alpha, torch.Tensor):
+            self.register_buffer("alpha", alpha.float())
+        elif alpha is None:
+            self.alpha = None
+        else:
+            raise ValueError(f"alpha must be float, Tensor, or None, got {type(alpha)}")
+    
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            logits: [B, C] unnormalized logits
+            targets: [B] class indices
+        """
+        ce_loss = F.cross_entropy(logits, targets, reduction="none")
+        pt = torch.exp(-ce_loss)  # p_t
+        focal_loss = ((1 - pt) ** self.gamma) * ce_loss
+        
+        if self.alpha is not None:
+            if self.alpha.dim() == 0 or len(self.alpha) == 1:
+                alpha_t = self.alpha.squeeze()
+            else:
+                alpha_t = self.alpha[targets]  # [B]
+            focal_loss = alpha_t * focal_loss
+        
+        if self.reduction == "mean":
+            return focal_loss.mean()
+        elif self.reduction == "none":
+            return focal_loss
+        else:
+            raise ValueError(f"reduction must be 'mean' or 'none', got {self.reduction}")
+
+
+class BalancedSoftmax(nn.Module):
+    """
+    Balanced Softmax: adjusted_logits = logits + log(class_counts)
+    Paper: https://arxiv.org/abs/2007.10740
+    
+    Args:
+        class_counts: [C] tensor of sample counts per class
+        reduction: 'mean' or 'none'
+    """
+    def __init__(self, class_counts: torch.Tensor, reduction: str = "mean"):
+        super().__init__()
+        if not isinstance(class_counts, torch.Tensor):
+            class_counts = torch.tensor(class_counts, dtype=torch.float32)
+        
+        class_counts = class_counts.float()
+        if (class_counts == 0).any():
+            zero_classes = (class_counts == 0).nonzero(as_tuple=True)[0].tolist()
+            raise ValueError(f"BalancedSoftmax requires non-zero class counts. Zero counts: {zero_classes}")
+        
+        self.register_buffer("log_class_counts", torch.log(class_counts))
+        self.reduction = reduction
+    
+    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            logits: [B, C] unnormalized logits
+            targets: [B] class indices
+        """
+        adjusted_logits = logits + self.log_class_counts.unsqueeze(0)
+        return F.cross_entropy(adjusted_logits, targets, reduction=self.reduction)

osf/models/base_pretrain_model.py

@@ -0,0 +1,144 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pytorch_lightning import LightningModule
+
+from osf.backbone.vit1d import vit_nano, vit_tiny, vit_small, vit_middle, vit_base
+
+VIT_FACTORIES = {
+    "vit_nano": vit_nano,
+    "vit_tiny": vit_tiny,
+    "vit_small": vit_small,
+    "vit_middle": vit_middle,
+    "vit_base": vit_base,
+}
+
+
+class PSGModalityEncoder(nn.Module):
+    """ViT encoder for PSG signals: backbone -> optional projection -> L2-norm"""
+
+    def __init__(self, *,
+                 encoder_name: str,
+                 proj_out: int = 256,
+                 proj_hidden: int = 512,
+                 freq: int = 64,
+                 win_sec: int = 30,
+                 channel: int = 11,
+                 lead_wise=0,
+                 patch_size=40,
+                 patch_size_ch=4,
+                 use_lead_embedding: bool = True,
+                 is_proj_head=1):
+        super().__init__()
+        token_len = freq * win_sec
+        self.token_len = token_len
+        self.patch_size = patch_size
+
+        if encoder_name not in VIT_FACTORIES:
+            raise ValueError(f"Unknown encoder_name: {encoder_name}. Choose from {list(VIT_FACTORIES.keys())}")
+
+        self.backbone = VIT_FACTORIES[encoder_name](
+            num_leads=channel, seq_len=token_len, patch_size=patch_size,
+            lead_wise=lead_wise, patch_size_ch=patch_size_ch,
+            use_lead_embedding=use_lead_embedding,
+        )
+
+        d_model = self.backbone.width
+        if is_proj_head == 1:
+            self.proj_head = nn.Sequential(
+                nn.Linear(d_model, proj_hidden),
+                nn.LayerNorm(proj_hidden),
+                nn.ReLU(inplace=True),
+                nn.Linear(proj_hidden, proj_out),
+                nn.LayerNorm(proj_out),
+            )
+        else:
+            self.proj_head = None
+
+    def forward(self, x, normalize=True):
+        # x: [B, C, T]
+        h = self.backbone(x)  # [B, D]
+        if self.proj_head is not None:
+            h = self.proj_head(h)  # [B, proj_out]
+        if normalize:
+            return F.normalize(h, dim=-1)
+        return h
+
+
+class BasePretrainModel(LightningModule):
+    def __init__(self,
+                 psg_encoder_name: str = "vit_base",
+                 text_encoder_name: str = "google/flan-t5-base",
+                 fusion_decoder_name: str = 'cross-attn',
+                 shared_emb_dim: int = 256,
+                 lr: float = 2e-4,
+                 weight_decay: float = 0.2,
+                 training_steps_per_epoch: int = 7000,
+                 max_epochs: int = 100,
+                 *args, **kwargs):
+        super().__init__()
+        self.save_hyperparameters()
+        self.psg_encoder_name = psg_encoder_name
+        self.text_encoder_name = text_encoder_name
+        self.fusion_decoder_name = fusion_decoder_name
+        self.shared_emb_dim = shared_emb_dim
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.training_steps_per_epoch = training_steps_per_epoch
+        self.max_epochs = max_epochs
+        self.warmup_epochs = 0.1 * self.max_epochs
+        self.proj_out = shared_emb_dim
+        self.proj_hidden = 256
+
+        assert self.training_steps_per_epoch > 1
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+            betas=(0.9, 0.95),
+        )
+
+        total_steps = int(self.training_steps_per_epoch * self.max_epochs)
+        warmup_steps = int(round(self.training_steps_per_epoch * self.warmup_epochs))
+        warmup_steps = max(0, warmup_steps)
+        decay_steps = max(1, total_steps - warmup_steps)
+
+        if warmup_steps > 0:
+            warmup = torch.optim.lr_scheduler.LinearLR(
+                optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_steps)
+            cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, T_max=decay_steps, eta_min=1e-8)
+            sched = torch.optim.lr_scheduler.SequentialLR(
+                optimizer, schedulers=[warmup, cosine], milestones=[warmup_steps])
+        else:
+            sched = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, T_max=decay_steps, eta_min=1e-8)
+
+        return [optimizer], [{"scheduler": sched, "interval": "step", "frequency": 1}]
+
+    def training_step(self, batch, batch_idx):
+        loss_dict, metrics_dict = self.shared_step(batch, batch_idx)
+        for k, v in loss_dict.items():
+            self.log(f"train/{k}", v, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
+        for k, v in metrics_dict.items():
+            self.log(f"train/{k}", v, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
+        return loss_dict['loss']
+
+    def validation_step(self, batch, batch_idx):
+        with torch.no_grad():
+            loss_dict, metrics_dict = self.shared_step(batch, batch_idx)
+        for k, v in loss_dict.items():
+            self.log(f"val/{k}", v, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        for k, v in metrics_dict.items():
+            self.log(f"val/{k}", v, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        return loss_dict
+
+    def test_step(self, batch, batch_idx):
+        loss_dict, metrics_dict = self.shared_step(batch, batch_idx)
+        for k, v in loss_dict.items():
+            self.log(f"test/{k}", v, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        for k, v in metrics_dict.items():
+            self.log(f"test/{k}", v, on_step=False, on_epoch=True, prog_bar=True, sync_dist=True)
+        return loss_dict

osf/models/base_pretrain_model_cls.py

@@ -0,0 +1,56 @@
+import torch.nn as nn
+from osf.backbone.vit1d_cls import vit_nano, vit_tiny, vit_small, vit_middle, vit_base, vit_large, vit_xl
+
+
+class PSGModalityEncoderCLS(nn.Module):
+    """
+    Init helper for ViT with CLS token. No forward() - access .backbone directly.
+    
+    Used by DINO to initialize encoder, then DINO accesses self.encoders["all"].backbone.
+    """
+    def __init__(self, *,
+                 encoder_name: str,
+                 proj_out: int      = 256,
+                 proj_hidden: int   = 512,
+                 freq: int          = 64,
+                 win_sec: int       = 30,
+                 channel: int       = 12, 
+                 lead_wise = 0,
+                 patch_size = 40, 
+                 patch_size_ch = 4,
+                 is_proj_head = 1,
+                ):
+        super().__init__()
+        token_len  = freq * win_sec
+
+        self.token_len = token_len
+        self.patch_size = patch_size
+
+        if encoder_name == "vit_nano":
+            self.backbone = vit_nano(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_tiny":
+            self.backbone = vit_tiny(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_small":
+            self.backbone = vit_small(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_middle":
+            self.backbone = vit_middle(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_base":
+            self.backbone = vit_base(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_large":
+            self.backbone = vit_large(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        elif encoder_name == "vit_xl":
+            self.backbone = vit_xl(num_leads=channel, seq_len=token_len, patch_size=patch_size, lead_wise=lead_wise, patch_size_ch=patch_size_ch)
+        else:
+            raise ValueError(f"Unknown encoder_name for CLS variant: {encoder_name}")
+
+        d_model = self.backbone.width
+        if is_proj_head == 1:
+            self.proj_head = nn.Sequential(
+                nn.Linear(d_model, proj_hidden),
+                nn.LayerNorm(proj_hidden),
+                nn.ReLU(inplace=True),
+                nn.Linear(proj_hidden, proj_out),
+                nn.LayerNorm(proj_out),
+            )
+        else:
+            self.proj_head = None

osf/models/dino_model_cls.py

@@ -0,0 +1,311 @@
+import copy
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+
+from osf.models.dino_utils.dino_clstoken_loss import DINOLoss
+from osf.models.dino_utils.ibot_patch_loss import iBOTPatchLoss
+from osf.models.dino_utils.koleo_loss import KoLeoLoss
+from osf.models.base_pretrain_model import BasePretrainModel
+from osf.models.base_pretrain_model_cls import PSGModalityEncoderCLS
+from osf.datasets.simclr_aug_registry import build_simclr_augmentor
+
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, hidden_dim=2048, bottleneck_dim=256, nlayers=3):
+        super().__init__()
+        num_layers = max(nlayers, 1)
+        if num_layers == 1:
+            self.mlp = nn.Sequential(nn.Linear(in_dim, bottleneck_dim))
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim), nn.GELU()]
+            for _ in range(num_layers - 2):
+                layers += [nn.Linear(hidden_dim, hidden_dim), nn.GELU()]
+            layers += [nn.Linear(hidden_dim, bottleneck_dim)]
+            self.mlp = nn.Sequential(*layers)
+
+        self.apply(self._init_weights)
+        self.prototypes = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.prototypes.weight_g.data.fill_(1.0)
+
+    @staticmethod
+    def _init_weights(m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = F.normalize(x, dim=-1)
+        return self.prototypes(x)
+
+
+class DINOCLSModel(BasePretrainModel):
+    def __init__(
+        self,
+        psg_encoder_name: str = "vit_base",
+        text_encoder_name: Optional[str] = None,
+        shared_emb_dim: int = 768,
+        out_dim: int = 2048,
+        patch_out_dim: int = 2048,
+        dino_out_dim: int = None,
+        dino_patch_out_dim: int = None,
+        dino_hidden_dim: int = 2048,
+        dino_bottleneck_dim: int = 256,
+        student_temp: float = 0.1,
+        teacher_temp_warmup: float = 0.04,
+        teacher_temp_final: float = 0.07,
+        teacher_temp_warmup_iters: int = 10000,
+        base_momentum: float = 0.996,
+        use_koleo: bool = True,
+        koleo_lambda: float = 0.0,
+        ibot_lambda: float = 0.0,
+        lr: float = 2e-4,
+        weight_decay: float = 0.2,
+        num_freeze_layers: int = 6,
+        simclr_augmentation: dict | None = None,
+        n_local_crops: int = 2,
+        *args, **kwargs
+    ):
+        super().__init__(
+            psg_encoder_name=psg_encoder_name,
+            text_encoder_name=None,
+            shared_emb_dim=shared_emb_dim,
+            lr=lr,
+            weight_decay=weight_decay,
+            *args, **kwargs
+        )
+        self.save_hyperparameters()
+
+        self.proj_out = shared_emb_dim
+        self.proj_hidden = 256
+        self.num_freeze_layers = num_freeze_layers
+
+        num_leads = kwargs.get('num_leads', 12)
+        self.num_leads = num_leads
+
+        self.cfg = [dict(name="all", freq=64, win_sec=30, in_ch=num_leads)]
+        self.encoders = nn.ModuleDict()
+        for mod in self.cfg:
+            self.encoders[mod["name"]] = PSGModalityEncoderCLS(
+                encoder_name=psg_encoder_name,
+                proj_out=shared_emb_dim,
+                proj_hidden=256,
+                freq=mod["freq"],
+                win_sec=mod["win_sec"],
+                channel=mod["in_ch"],
+                patch_size=kwargs['patch_size_time'],
+                lead_wise=kwargs['lead_wise'],
+                patch_size_ch=(num_leads if kwargs['lead_wise'] == 0 else kwargs['patch_size_ch']),
+                is_proj_head=0,
+            )
+        self.lead_wise = kwargs['lead_wise']
+        self.patch_size_time = kwargs['patch_size_time']
+        self.patch_size_ch = (num_leads if self.lead_wise == 0 else kwargs['patch_size_ch'])
+        trunk_dim = self.encoders['all'].backbone.width
+        out_dim = dino_out_dim if dino_out_dim is not None else out_dim
+        patch_out_dim = dino_patch_out_dim if dino_patch_out_dim is not None else patch_out_dim
+        self.out_dim = out_dim
+        self.patch_out_dim = patch_out_dim
+
+        self.student_global_head = DINOHead(trunk_dim, out_dim, dino_hidden_dim, dino_bottleneck_dim, 3)
+        self.student_patch_head = DINOHead(trunk_dim, patch_out_dim, dino_hidden_dim, dino_bottleneck_dim, 3)
+        self.teacher_encoder = copy.deepcopy(self.encoders["all"])
+        for p in self.teacher_encoder.parameters():
+            p.requires_grad = False
+        self.teacher_encoder.eval()
+
+        self.teacher_global_head = DINOHead(trunk_dim, out_dim, dino_hidden_dim, dino_bottleneck_dim, 3)
+        self.teacher_patch_head = DINOHead(trunk_dim, patch_out_dim, dino_hidden_dim, dino_bottleneck_dim, 3)
+        self.teacher_global_head.load_state_dict(self.student_global_head.state_dict(), strict=True)
+        self.teacher_patch_head.load_state_dict(self.student_patch_head.state_dict(), strict=True)
+        for p in self.teacher_global_head.parameters():
+            p.requires_grad = False
+        for p in self.teacher_patch_head.parameters():
+            p.requires_grad = False
+        self.teacher_global_head.eval()
+        self.teacher_patch_head.eval()
+        self.dino_loss = DINOLoss(out_dim=out_dim, student_temp=student_temp, center_momentum=0.9)
+        self.ibot_loss = iBOTPatchLoss(patch_out_dim=patch_out_dim, student_temp=student_temp, center_momentum=0.9)
+        self.koleo = KoLeoLoss() if use_koleo else None
+        self.koleo_lambda = float(koleo_lambda)
+        self.ibot_lambda = float(ibot_lambda)
+        self.teacher_temp_warmup = float(teacher_temp_warmup)
+        self.teacher_temp_final = float(teacher_temp_final)
+        self.teacher_temp_warmup_iters = int(teacher_temp_warmup_iters)
+        self.base_momentum = float(base_momentum)
+
+        self.register_buffer("seen_steps", torch.tensor(0, dtype=torch.long))
+
+        if simclr_augmentation is None:
+            simclr_augmentation = {}
+        self.simclr_augmentation = simclr_augmentation
+        self.augmentor = build_simclr_augmentor(self.simclr_augmentation)
+        self.n_local_crops = int(n_local_crops)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, trunk_dim))
+        nn.init.trunc_normal_(self.mask_token, std=0.02)
+
+    def _teacher_temp(self, step: int) -> float:
+        if step < self.teacher_temp_warmup_iters:
+            alpha = step / float(max(1, self.teacher_temp_warmup_iters))
+            return self.teacher_temp_warmup * (1 - alpha) + self.teacher_temp_final * alpha
+        return self.teacher_temp_final
+
+    def _momentum(self, step: int, max_steps: int) -> float:
+        return 1.0 - (1.0 - self.base_momentum) * (math.cos(math.pi * step / max_steps) + 1) / 2
+
+    @torch.no_grad()
+    def _ema_update(self, m: float):
+        for param_q, param_k in zip(self.encoders['all'].parameters(), self.teacher_encoder.parameters()):
+            param_k.data.mul_(m).add_(param_q.data, alpha=1.0 - m)
+        for param_q, param_k in zip(self.student_global_head.parameters(), self.teacher_global_head.parameters()):
+            param_k.data.mul_(m).add_(param_q.data, alpha=1.0 - m)
+        for param_q, param_k in zip(self.student_patch_head.parameters(), self.teacher_patch_head.parameters()):
+            param_k.data.mul_(m).add_(param_q.data, alpha=1.0 - m)
+        self.teacher_encoder.eval()
+        self.teacher_global_head.eval()
+        self.teacher_patch_head.eval()
+
+    def _forward_encoder(self, encoder, x, return_tokens=True):
+        # x: [B, C, T]
+        if return_tokens:
+            cls, patches = encoder.backbone.forward_encoding(x, return_sequence=False)
+            return cls, patches  # [B, D], [B, N, D]
+        else:
+            cls = encoder.backbone(x)
+            return cls, None  # [B, D], None
+
+    def _make_views_aug(self, x: torch.Tensor):
+        v1, v2 = self.augmentor(x)
+        globals_x = [v1, v2]
+        locals_x = []
+        for _ in range(self.n_local_crops):
+            lv1, _ = self.augmentor(x)
+            locals_x.append(lv1)
+        return globals_x, locals_x
+
+    def shared_step(self, batch, batch_idx):
+        x = batch["psg"]
+        globals_x, locals_x = self._make_views_aug(x)
+        tt = self._teacher_temp(int(self.global_step))
+
+        with torch.no_grad():
+            teacher_out_soft_list = []
+            teacher_global_logits_cache = []
+            teacher_patch_logits_cache = []
+
+            if len(globals_x) > 0:
+                g_sizes = [gx.size(0) for gx in globals_x]
+                g_cat = torch.cat(globals_x, dim=0)
+                cls_t_cat, _ = self._forward_encoder(self.teacher_encoder, g_cat, return_tokens=True)
+                g_logits_cat = self.teacher_global_head(cls_t_cat)
+                g_logits_split = list(torch.split(g_logits_cat, g_sizes, dim=0))
+                teacher_out_soft_list = [self.dino_loss.softmax_center_teacher(gl, tt) for gl in g_logits_split]
+                teacher_global_logits_cache = g_logits_split
+
+        student_global_logits = []
+        student_cls_tokens = []
+        all_student_views = globals_x + locals_x
+        if len(all_student_views) > 0:
+            s_sizes = [sx.size(0) for sx in all_student_views]
+            s_cat = torch.cat(all_student_views, dim=0)
+            cls_s_cat, _ = self._forward_encoder(self.encoders["all"], s_cat, return_tokens=False)
+            sg_logits_cat = self.student_global_head(cls_s_cat)
+            student_global_logits = list(torch.split(sg_logits_cat, s_sizes, dim=0))
+            student_cls_tokens = list(torch.split(cls_s_cat, s_sizes, dim=0))
+
+        ibot_loss_val = torch.tensor(0.0, device=x.device)
+        if len(globals_x) > 0:
+            with torch.no_grad():
+                t_tokens, _ = self.teacher_encoder.backbone.to_tokens_2d(
+                    globals_x[0], patch_size_ch=self.patch_size_ch, patch_size_time=self.patch_size_time)
+                B2 = t_tokens.size(0)
+                cls_tok = self.teacher_encoder.backbone.cls_token.expand(B2, -1, -1)
+                t_full = torch.cat([cls_tok, t_tokens], dim=1)
+                pe_full = self.teacher_encoder.backbone.pos_embedding[:, :t_full.size(1), :].to(t_full.device)
+                t_full = t_full + pe_full
+                t_full = self.teacher_encoder.backbone._run_blocks(t_full)
+                _, t_patches = t_full[:, 0], t_full[:, 1:]
+                t_logits_all = self.teacher_patch_head(t_patches)
+                t_soft = self.ibot_loss.softmax_center_teacher(t_logits_all, tt)
+
+            s_tokens, _ = self.encoders["all"].backbone.to_tokens_2d(
+                globals_x[0], patch_size_ch=self.patch_size_ch, patch_size_time=self.patch_size_time)
+            B2, N, Dtok = s_tokens.shape
+
+            mask_ratio = float(getattr(self, "ibot_mask_ratio", 0.3))
+            n_mask = max(1, int(round(N * mask_ratio)))
+            rand = torch.rand(B2, N, device=x.device)
+            topk_idx = rand.topk(k=n_mask, dim=1, largest=True).indices
+            masks = torch.zeros(B2, N, dtype=torch.bool, device=x.device)
+            masks.scatter_(1, topk_idx, True)
+
+            s_tokens_masked = torch.where(
+                masks.unsqueeze(-1),
+                self.mask_token.expand_as(s_tokens),
+                s_tokens
+            )
+
+            cls_tok_s = self.encoders["all"].backbone.cls_token.expand(B2, -1, -1)
+            s_full = torch.cat([cls_tok_s, s_tokens_masked], dim=1)
+            pe_full_s = self.encoders["all"].backbone.pos_embedding[:, :s_full.size(1), :].to(s_full.device)
+            s_full = s_full + pe_full_s
+            s_full = self.encoders["all"].backbone._run_blocks(s_full)
+            _, s_patches = s_full[:, 0], s_full[:, 1:]
+            s_logits_all = self.student_patch_head(s_patches)
+
+            ibot_loss_val = self.ibot_loss.forward_masked(
+                student_patch_tokens_masked=s_logits_all[masks],
+                teacher_patch_tokens_masked=t_soft[masks],
+                student_masks_flat=masks,
+            )
+
+            with torch.no_grad():
+                teacher_patch_logits_cache.append(t_logits_all)
+
+        dino_loss_val = self.dino_loss(student_global_logits, teacher_out_soft_list)
+        pair_norm = max(1, len(student_global_logits) * len(teacher_out_soft_list))
+        dino_loss_val = dino_loss_val / pair_norm
+        koleo_val = torch.tensor(0.0, device=x.device)
+        if self.koleo is not None and len(student_cls_tokens) > 0:
+            koleo_val = self.koleo(F.normalize(student_cls_tokens[0], dim=-1))
+
+        total_loss = dino_loss_val + self.ibot_lambda * ibot_loss_val + self.koleo_lambda * koleo_val
+
+        with torch.no_grad():
+            if self.training:
+                if len(teacher_global_logits_cache) > 0:
+                    self.dino_loss.update_center(torch.cat(teacher_global_logits_cache, dim=0))
+                if len(teacher_patch_logits_cache) > 0:
+                    self.ibot_loss.update_center(torch.cat(teacher_patch_logits_cache, dim=0))
+
+        metrics = {
+            "loss": total_loss,
+            "loss/dino": dino_loss_val,
+            "loss/ibot": ibot_loss_val,
+            "loss/koleo": koleo_val,
+            "sched/teacher_temp": torch.tensor(tt, device=x.device),
+        }
+        return {"loss": total_loss}, metrics
+
+    def training_step(self, batch, batch_idx):
+        loss_dict, metrics = self.shared_step(batch, batch_idx)
+        for k, v in metrics.items():
+            self.log(f"train/{k}", v, on_step=True, on_epoch=True, prog_bar=(k == "loss"), sync_dist=True)
+        return loss_dict["loss"]
+
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        max_steps = max(1, getattr(self.trainer, "max_steps", getattr(self.trainer, "estimated_stepping_batches", 100000)))
+        m = self._momentum(int(self.global_step), max_steps)
+        self._ema_update(m)
+        self.log("sched/momentum", torch.tensor(m, device=self.device), on_step=True, prog_bar=False)
+
+    def validation_step(self, batch, batch_idx):
+        loss_dict, metrics = self.shared_step(batch, batch_idx)
+        for k, v in metrics.items():
+            self.log(f"val/{k}", v, on_step=True, on_epoch=True, prog_bar=(k == "loss"), sync_dist=True)
+        return loss_dict["loss"]

osf/models/dino_utils/dino_clstoken_loss.py

@@ -0,0 +1,96 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+
+class DINOLoss(nn.Module):
+    def __init__(
+        self,
+        out_dim,
+        student_temp=0.1,
+        center_momentum=0.9,
+    ):
+        super().__init__()
+        self.student_temp = student_temp
+        self.center_momentum = center_momentum
+        self.register_buffer("center", torch.zeros(1, out_dim))
+        self.updated = True
+        self.reduce_handle = None
+        self.len_teacher_output = None
+        self.async_batch_center = None
+
+    @torch.no_grad()
+    def softmax_center_teacher(self, teacher_output, teacher_temp):
+        self.apply_center_update()
+        # teacher centering and sharpening
+        return F.softmax((teacher_output - self.center) / teacher_temp, dim=-1)
+
+    @torch.no_grad()
+    def sinkhorn_knopp_teacher(self, teacher_output, teacher_temp, n_iterations=3):
+        teacher_output = teacher_output.float()
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        Q = torch.exp(teacher_output / teacher_temp).t()
+        B = Q.shape[1] * world_size
+        K = Q.shape[0]
+        sum_Q = torch.sum(Q)
+        if dist.is_initialized():
+            dist.all_reduce(sum_Q)
+        Q /= sum_Q
+
+        for it in range(n_iterations):
+            sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
+            if dist.is_initialized():
+                dist.all_reduce(sum_of_rows)
+            Q /= sum_of_rows
+            Q /= K
+            Q /= torch.sum(Q, dim=0, keepdim=True)
+            Q /= B
+
+        Q *= B
+        return Q.t()
+
+    def forward(self, student_output_list, teacher_out_softmaxed_centered_list):
+        """
+        Cross-entropy between softmax outputs of the teacher and student networks.
+        """
+        # TODO: Use cross_entropy_distribution here
+        total_loss = 0
+        for s in student_output_list:
+            lsm = F.log_softmax(s / self.student_temp, dim=-1)
+            for t in teacher_out_softmaxed_centered_list:
+                loss = torch.sum(t * lsm, dim=-1)
+                total_loss -= loss.mean()
+        return total_loss
+
+    @torch.no_grad()
+    def update_center(self, teacher_output):
+        self.reduce_center_update(teacher_output)
+
+    @torch.no_grad()
+    def reduce_center_update(self, teacher_output):
+        self.updated = False
+        self.len_teacher_output = len(teacher_output)
+        self.async_batch_center = torch.sum(teacher_output, dim=0, keepdim=True)
+        if dist.is_initialized():
+            self.reduce_handle = dist.all_reduce(self.async_batch_center, async_op=True)
+
+    @torch.no_grad()
+    def apply_center_update(self):
+        if self.updated is False:
+            world_size = dist.get_world_size() if dist.is_initialized() else 1
+
+            if self.reduce_handle is not None:
+                self.reduce_handle.wait()
+            _t = self.async_batch_center / (self.len_teacher_output * world_size)
+
+            self.center = self.center * self.center_momentum + _t * (1 - self.center_momentum)
+
+            self.updated = True
+
+

osf/models/dino_utils/ibot_patch_loss.py

@@ -0,0 +1,134 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch import nn
+
+import logging
+
+
+logger = logging.getLogger("dinov2")
+
+
+try:
+    from xformers.ops import cross_entropy
+
+    def lossfunc(t, s, temp):
+        s = s.float()
+        t = t.float()
+        if s.ndim == 2:
+            return -cross_entropy(s.unsqueeze(0), t.unsqueeze(0), temp, bw_inplace=True).squeeze(0)
+        elif s.ndim == 3:
+            return -cross_entropy(s, t, temp, bw_inplace=True)
+
+except ImportError:
+
+    def lossfunc(t, s, temp):
+        return torch.sum(t * F.log_softmax(s / temp, dim=-1), dim=-1)
+
+
+class iBOTPatchLoss(nn.Module):
+    def __init__(self, patch_out_dim, student_temp=0.1, center_momentum=0.9):
+        super().__init__()
+        self.student_temp = student_temp
+        self.center_momentum = center_momentum
+        self.register_buffer("center", torch.zeros(1, 1, patch_out_dim))
+        self.updated = True
+        self.reduce_handle = None
+        self.len_teacher_patch_tokens = None
+        self.async_batch_center = None
+
+    @torch.no_grad()
+    def softmax_center_teacher(self, teacher_patch_tokens, teacher_temp):
+        self.apply_center_update()
+        return F.softmax((teacher_patch_tokens - self.center) / teacher_temp, dim=-1)
+
+    @torch.no_grad()
+    def sinkhorn_knopp_teacher(self, teacher_output, teacher_temp, n_masked_patches_tensor, n_iterations=3):
+        teacher_output = teacher_output.float()
+        # world_size = dist.get_world_size() if dist.is_initialized() else 1
+        Q = torch.exp(teacher_output / teacher_temp).t()
+        B = n_masked_patches_tensor
+        dist.all_reduce(B)
+        K = Q.shape[0]
+        sum_Q = torch.sum(Q)
+        if dist.is_initialized():
+            dist.all_reduce(sum_Q)
+        Q /= sum_Q
+
+        for it in range(n_iterations):
+            sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
+            if dist.is_initialized():
+                dist.all_reduce(sum_of_rows)
+            Q /= sum_of_rows
+            Q /= K
+            Q /= torch.sum(Q, dim=0, keepdim=True)
+            Q /= B
+
+        Q *= B
+        return Q.t()
+
+    def forward(self, student_patch_tokens, teacher_patch_tokens, student_masks_flat):
+        """
+        Cross-entropy between softmax outputs of the teacher and student networks.
+        student_patch_tokens: (B, N, D) tensor
+        teacher_patch_tokens: (B, N, D) tensor
+        student_masks_flat: (B, N) tensor
+        """
+        t = teacher_patch_tokens
+        s = student_patch_tokens
+        loss = torch.sum(t * F.log_softmax(s / self.student_temp, dim=-1), dim=-1)
+        loss = torch.sum(loss * student_masks_flat.float(), dim=-1) / student_masks_flat.sum(dim=-1).clamp(min=1.0)
+        return -loss.mean()
+
+    def forward_masked(
+        self,
+        student_patch_tokens_masked,
+        teacher_patch_tokens_masked,
+        student_masks_flat,
+        n_masked_patches=None,
+        masks_weight=None,
+    ):
+        t = teacher_patch_tokens_masked
+        s = student_patch_tokens_masked
+        # loss = torch.sum(t * F.log_softmax(s / self.student_temp, dim=-1), dim=-1)
+        loss = lossfunc(t, s, self.student_temp)
+        if masks_weight is None:
+            masks_weight = (
+                (1 / student_masks_flat.sum(-1).clamp(min=1.0))
+                .unsqueeze(-1)
+                .expand_as(student_masks_flat)[student_masks_flat]
+            )
+        if n_masked_patches is not None:
+            loss = loss[:n_masked_patches]
+        loss = loss * masks_weight
+        return -loss.sum() / student_masks_flat.shape[0]
+
+    @torch.no_grad()
+    def update_center(self, teacher_patch_tokens):
+        self.reduce_center_update(teacher_patch_tokens)
+
+    @torch.no_grad()
+    def reduce_center_update(self, teacher_patch_tokens):
+        self.updated = False
+        self.len_teacher_patch_tokens = len(teacher_patch_tokens)
+        self.async_batch_center = torch.sum(teacher_patch_tokens.mean(1), dim=0, keepdim=True)
+        if dist.is_initialized():
+            self.reduce_handle = dist.all_reduce(self.async_batch_center, async_op=True)
+
+    @torch.no_grad()
+    def apply_center_update(self):
+        if self.updated is False:
+            world_size = dist.get_world_size() if dist.is_initialized() else 1
+
+            if self.reduce_handle is not None:
+                self.reduce_handle.wait()
+            _t = self.async_batch_center / (self.len_teacher_patch_tokens * world_size)
+
+            self.center = self.center * self.center_momentum + _t * (1 - self.center_momentum)
+
+            self.updated = True

osf/models/dino_utils/koleo_loss.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# import torch.distributed as dist
+
+
+logger = logging.getLogger("dinov2")
+
+
+class KoLeoLoss(nn.Module):
+    """Kozachenko-Leonenko entropic loss regularizer from Sablayrolles et al. - 2018 - Spreading vectors for similarity search"""
+
+    def __init__(self):
+        super().__init__()
+        self.pdist = nn.PairwiseDistance(2, eps=1e-8)
+
+    def pairwise_NNs_inner(self, x):
+        """
+        Pairwise nearest neighbors for L2-normalized vectors.
+        Uses Torch rather than Faiss to remain on GPU.
+        """
+        dots = torch.mm(x, x.t())
+        n = x.shape[0]
+        dots.view(-1)[:: (n + 1)].fill_(-1)
+        _, I = torch.max(dots, dim=1)  # noqa: E741
+        return I
+
+    def forward(self, student_output, eps=1e-8):
+        """
+        Args:
+            student_output (BxD): backbone output of student
+        """
+        with torch.cuda.amp.autocast(enabled=False):
+            student_output = F.normalize(student_output, eps=eps, p=2, dim=-1)
+            I = self.pairwise_NNs_inner(student_output)  # noqa: E741
+            distances = self.pdist(student_output, student_output[I])  # BxD, BxD -> B
+            loss = -torch.log(distances + eps).mean()
+        return loss

osf/models/ssl_finetuner.py

@@ -0,0 +1,568 @@
+from typing import Tuple, Optional
+
+import torch
+import torch.nn as nn
+from pytorch_lightning import LightningModule
+import torch.nn.functional as F
+from einops import rearrange
+from itertools import chain
+from torchmetrics import Accuracy, Precision, Recall, F1Score, AUROC, ConfusionMatrix, CohenKappa, AveragePrecision, MetricCollection
+from osf.models.balanced_losses import FocalLoss, BalancedSoftmax
+
+
+def _create_pred_metrics(num_classes: int) -> MetricCollection:
+    """Create metrics that take preds (class indices) as input."""
+    metrics = {
+        "acc": Accuracy(task="multiclass", num_classes=num_classes, average="micro"),
+        "f1": F1Score(task="multiclass", num_classes=num_classes, average="macro"),
+        "f1_w": F1Score(task="multiclass", num_classes=num_classes, average="weighted"),
+        "rec_m": Recall(task="multiclass", num_classes=num_classes, average="macro"),
+        "kappa": CohenKappa(task="multiclass", num_classes=num_classes, weights="quadratic"),
+    }
+    return MetricCollection(metrics)
+
+
+def _create_prob_metrics(num_classes: int) -> MetricCollection:
+    """Create metrics that take probs (probabilities) as input."""
+    metrics = {
+        "auc": AUROC(task="multiclass", num_classes=num_classes, average="macro"),
+        "auprc": AveragePrecision(task="multiclass", num_classes=num_classes, average="macro"),
+    }
+    return MetricCollection(metrics)
+
+
+def _create_perclass_pred_metrics(num_classes: int) -> MetricCollection:
+    """Create per-class metrics that take preds as input."""
+    metrics = {
+        "acc_c": Accuracy(task="multiclass", num_classes=num_classes, average=None),
+        "prec_c": Precision(task="multiclass", num_classes=num_classes, average=None),
+        "rec_c": Recall(task="multiclass", num_classes=num_classes, average=None),
+        "f1_c": F1Score(task="multiclass", num_classes=num_classes, average=None),
+        "cm": ConfusionMatrix(task="multiclass", num_classes=num_classes, normalize=None),
+    }
+    return MetricCollection(metrics)
+
+
+def _create_perclass_prob_metrics(num_classes: int) -> MetricCollection:
+    """Create per-class metrics that take probs as input."""
+    metrics = {
+        "auc_c": AUROC(task="multiclass", num_classes=num_classes, average=None),
+        "auprc_c": AveragePrecision(task="multiclass", num_classes=num_classes, average=None),
+    }
+    return MetricCollection(metrics)
+
+
+
+class SSLFineTuner(LightningModule):
+    def __init__(self,
+        backbones,
+        use_which_backbone,
+        config = None,
+        in_features: int = 256,
+        num_classes: int = 2,
+        epochs: int = 10,
+        dropout: float = 0.0,
+        lr: float = 1e-3,
+        weight_decay: float = 1e-4,
+        final_lr: float = 1e-5,
+        use_channel_bank: bool = True,
+        loss_type: str = "ce",
+        class_distribution: Optional[torch.Tensor] = None,
+        focal_gamma: float = 2.0,
+        focal_alpha: Optional[float | torch.Tensor] = None,
+        use_mean_pool: bool = False,
+        total_training_steps: int = None,
+        finetune_backbone: bool = False,
+        *args, **kwargs
+    ) -> None:
+        super().__init__()
+        self.save_hyperparameters()
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.epochs = epochs
+        self.final_lr = final_lr
+        self.use_channel_bank = use_channel_bank
+        self.loss_type = loss_type
+        self.focal_gamma = focal_gamma
+        self.focal_alpha = focal_alpha
+        self.use_mean_pool = use_mean_pool
+        self.total_training_steps = total_training_steps
+        self.finetune_backbone = finetune_backbone
+        
+        if loss_type == "ce":
+            self.criterion = None 
+        elif loss_type == "focal":
+            alpha = focal_alpha
+            if alpha is None and class_distribution is not None:
+                class_dist = class_distribution.float()
+                total_samples = class_dist.sum()
+                alpha = total_samples / (num_classes * class_dist)
+                alpha = alpha / alpha.mean()
+            self.criterion = FocalLoss(alpha=alpha, gamma=focal_gamma, reduction="mean")
+        elif loss_type == "balanced_softmax":
+            self.criterion = BalancedSoftmax(class_distribution, reduction="mean")
+        else:
+            raise ValueError(f"Unknown loss_type: {loss_type}. Must be one of ['ce', 'focal', 'balanced_softmax']")
+
+        if isinstance(backbones, nn.ModuleDict):
+            self.backbones = backbones
+        else:
+            self.backbones = nn.ModuleDict(backbones)
+        self.config = config
+        self.use_which_backbone = use_which_backbone
+        self.backbone = self.backbones[self.use_which_backbone] if self.use_which_backbone != "fusion" else None
+
+
+        if self.use_which_backbone == "fusion":
+            for k in ("ecg", "resp", "elect"):
+                if k in self.backbones:
+                    for p in self.backbones[k].parameters():
+                        p.requires_grad = self.finetune_backbone
+                    if not self.finetune_backbone:
+                        self.backbones[k].eval()
+        else:
+            for p in self.backbone.parameters():
+                p.requires_grad = self.finetune_backbone
+            if not self.finetune_backbone:
+                self.backbone.eval()
+        
+        if self.finetune_backbone:
+            print(f"[INFO] Full finetuning mode: backbone parameters are TRAINABLE")
+
+        if self.use_which_backbone == "fusion":
+            dims = [getattr(self.backbones[k], "out_dim", in_features)
+                    for k in ("ecg", "resp", "elect") if k in self.backbones]
+            if len(dims) == 0:
+                raise ValueError("fusion requires at least one of {'ecg','resp','elect'} in backbones.")
+            if len(set(dims)) != 1:
+                raise ValueError(f"Mean fusion requires equal output dims, got {dims}")
+            final_in_features = dims[0]
+        else:
+            final_in_features = getattr(self.backbone, "out_dim", in_features)
+
+        self.linear_layer = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(final_in_features, num_classes)
+        )
+
+        self.train_pred_metrics = _create_pred_metrics(num_classes)
+        self.val_pred_metrics = _create_pred_metrics(num_classes)
+        self.test_pred_metrics = _create_pred_metrics(num_classes)
+
+        self.train_prob_metrics = _create_prob_metrics(num_classes)
+        self.val_prob_metrics = _create_prob_metrics(num_classes)
+        self.test_prob_metrics = _create_prob_metrics(num_classes)
+
+        self.train_pred_metrics_c = _create_perclass_pred_metrics(num_classes)
+        self.val_pred_metrics_c = _create_perclass_pred_metrics(num_classes)
+        self.test_pred_metrics_c = _create_perclass_pred_metrics(num_classes)
+
+        self.train_prob_metrics_c = _create_perclass_prob_metrics(num_classes)
+        self.val_prob_metrics_c = _create_perclass_prob_metrics(num_classes)
+        self.test_prob_metrics_c = _create_perclass_prob_metrics(num_classes)
+
+
+        self.class_names = getattr(self.config, "class_names", [str(i) for i in range(num_classes)])
+
+    def on_train_epoch_start(self) -> None:
+        if not self.finetune_backbone:
+            if self.use_which_backbone == "fusion":
+                for k in ("ecg", "resp", "elect"):
+                    if k in self.backbones:
+                        self.backbones[k].eval()
+            else:
+                self.backbone.eval()
+
+    def training_step(self, batch, batch_idx):
+        loss, logits, y = self.shared_step(batch)
+        probs = logits.softmax(-1)
+        preds = logits.argmax(-1)
+
+        self.train_pred_metrics.update(preds, y)
+        self.train_prob_metrics.update(probs, y)
+        self.train_pred_metrics_c.update(preds, y)
+        self.train_prob_metrics_c.update(probs, y)
+
+        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True)
+        return loss
+
+    def on_train_epoch_end(self):
+        pred_agg = self.train_pred_metrics.compute()
+        prob_agg = self.train_prob_metrics.compute()
+        
+        self.log("train_acc", pred_agg["acc"], prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("train_f1", pred_agg["f1"], prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("train_auc", prob_agg["auc"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("train_auprc", prob_agg["auprc"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+
+        pred_c = self.train_pred_metrics_c.compute()
+        prob_c = self.train_prob_metrics_c.compute()
+        cm = pred_c["cm"]
+        support = cm.sum(dim=1) if cm is not None else None
+
+        for i in range(len(pred_c["acc_c"])):
+            name = self.class_names[i] if i < len(self.class_names) else str(i)
+            self.log(f"train/acc_{name}", pred_c["acc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"train/prec_{name}", pred_c["prec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"train/rec_{name}", pred_c["rec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"train/f1_{name}", pred_c["f1_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"train/auc_{name}", prob_c["auc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"train/auprc_{name}", prob_c["auprc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            if support is not None:
+                self.log(f"train/support_{name}", support[i].to(pred_c["acc_c"][i].dtype), on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+
+        self.train_pred_metrics.reset()
+        self.train_prob_metrics.reset()
+        self.train_pred_metrics_c.reset()
+        self.train_prob_metrics_c.reset()
+
+    def validation_step(self, batch, batch_idx):
+        loss, logits, y = self.shared_step(batch)
+        probs = logits.softmax(-1)
+        preds = logits.argmax(-1)
+
+        self.val_pred_metrics.update(preds, y)
+        self.val_prob_metrics.update(probs, y)
+        self.val_pred_metrics_c.update(preds, y)
+        self.val_prob_metrics_c.update(probs, y)
+
+        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        return loss
+
+    def on_validation_epoch_end(self):
+        pred_agg = self.val_pred_metrics.compute()
+        prob_agg = self.val_prob_metrics.compute()
+        
+        self.log("val_acc", pred_agg["acc"], prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_f1", pred_agg["f1"], prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_f1_w", pred_agg["f1_w"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_rec_m", pred_agg["rec_m"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_auc", prob_agg["auc"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_auprc", prob_agg["auprc"], prog_bar=False, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_kappa", pred_agg["kappa"], prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+
+        pred_c = self.val_pred_metrics_c.compute()
+        prob_c = self.val_prob_metrics_c.compute()
+        cm = pred_c["cm"]
+        support = cm.sum(dim=1)
+
+        for i in range(len(pred_c["acc_c"])):
+            name = self.class_names[i] if i < len(self.class_names) else str(i)
+            self.log(f"val/acc_{name}", pred_c["acc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/prec_{name}", pred_c["prec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/rec_{name}", pred_c["rec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/f1_{name}", pred_c["f1_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/auc_{name}", prob_c["auc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/auprc_{name}", prob_c["auprc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"val/support_{name}", support[i].to(pred_c["acc_c"][i].dtype), on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+
+        self.val_pred_metrics.reset()
+        self.val_prob_metrics.reset()
+        self.val_pred_metrics_c.reset()
+        self.val_prob_metrics_c.reset()
+
+    def test_step(self, batch, batch_idx):
+        loss, logits, y = self.shared_step(batch)
+        probs = logits.softmax(-1)
+        preds = logits.argmax(-1)
+
+        self.test_pred_metrics.update(preds, y)
+        self.test_prob_metrics.update(probs, y)
+        self.test_pred_metrics_c.update(preds, y)
+        self.test_prob_metrics_c.update(probs, y)
+
+        self.log("test_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
+        return loss
+
+    def on_test_epoch_end(self):
+        pred_agg = self.test_pred_metrics.compute()
+        prob_agg = self.test_prob_metrics.compute()
+        
+        self.log("test_acc", pred_agg["acc"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_f1", pred_agg["f1"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_f1_w", pred_agg["f1_w"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_rec_m", pred_agg["rec_m"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_auc", prob_agg["auc"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_auprc", prob_agg["auprc"], on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_kappa", pred_agg["kappa"], on_step=False, on_epoch=True, sync_dist=True)
+
+        pred_c = self.test_pred_metrics_c.compute()
+        prob_c = self.test_prob_metrics_c.compute()
+        cm = pred_c["cm"]
+        support = cm.sum(dim=1) if cm is not None else None
+
+        for i in range(len(pred_c["acc_c"])):
+            name = self.class_names[i] if i < len(self.class_names) else str(i)
+            self.log(f"test/acc_{name}", pred_c["acc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"test/prec_{name}", pred_c["prec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"test/rec_{name}", pred_c["rec_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"test/f1_{name}", pred_c["f1_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"test/auc_{name}", prob_c["auc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            self.log(f"test/auprc_{name}", prob_c["auprc_c"][i], on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+            if support is not None:
+                self.log(f"test/support_{name}", support[i].to(pred_c["acc_c"][i].dtype),
+                         on_step=False, on_epoch=True, prog_bar=False, sync_dist=True)
+
+        self.test_pred_metrics.reset()
+        self.test_prob_metrics.reset()
+        self.test_pred_metrics_c.reset()
+        self.test_prob_metrics_c.reset()
+    def shared_step(self, batch):
+        context = torch.no_grad() if not self.finetune_backbone else torch.enable_grad()
+        
+        with context:
+            psg = batch['psg']
+            if self.use_which_backbone == 'ecg':
+                x = psg[:, 0:1, :]
+                feats = self._get_features(self.backbone, x)
+                
+            elif self.use_which_backbone == 'resp':
+                x = psg[:, 1:5, :]
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'elect':
+                x = psg[:, 5:, :]
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'all':
+
+                x = psg
+                feats = self._get_features(self.backbone, x)
+                
+            elif self.use_which_backbone == 'fusion':
+                feats_list = []
+                if 'ecg' in self.backbones:
+                    x_ecg = psg[:, 0:1, :]
+                    f_ecg = self._get_features(self.backbones['ecg'], x_ecg)
+                    feats_list.append(f_ecg)
+                if 'resp' in self.backbones:
+                    x_resp = psg[:, 1:5, :]
+                    f_resp = self._get_features(self.backbones['resp'], x_resp)
+                    feats_list.append(f_resp)
+                if 'elect' in self.backbones:
+                    x_elect = psg[:, 5:, :]
+                    f_elect = self._get_features(self.backbones['elect'], x_elect)
+                    feats_list.append(f_elect)
+
+
+                feats = torch.stack(feats_list, dim=0).mean(dim=0)
+            else:
+                raise ValueError(f"Unknown use_which_backbone: {self.use_which_backbone}")
+
+        y = batch["label"]
+        feats = feats.view(feats.size(0), -1)
+        logits = self.linear_layer(feats)
+        y = y.squeeze(1).long()
+
+        if self.criterion is None:
+            loss = F.cross_entropy(logits, y)
+        else:
+            loss = self.criterion(logits, y)
+        
+        return loss, logits, y
+    
+    def _get_features(self, backbone, x):
+        """Get features from backbone. Uses mean pooling if use_mean_pool=True."""
+        if self.use_mean_pool:
+            if hasattr(backbone, 'forward_encoding_mean_pool'):
+                return backbone.forward_encoding_mean_pool(x)
+            elif hasattr(backbone, 'forward_avg_pool'):
+                return backbone.forward_avg_pool(x)
+        return backbone(x)
+
+
+    def configure_optimizers(self):
+        if self.finetune_backbone:
+            if self.use_which_backbone == "fusion":
+                backbone_params = chain(*[self.backbones[k].parameters() 
+                                          for k in ("ecg", "resp", "elect") if k in self.backbones])
+            else:
+                backbone_params = self.backbone.parameters()
+            params = chain(backbone_params, self.linear_layer.parameters())
+        else:
+            params = self.linear_layer.parameters()
+        
+        optimizer = torch.optim.AdamW(
+            params,
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+        )
+
+        if self.total_training_steps is not None and self.total_training_steps > 0:
+            warmup_steps = int(0.1 * self.total_training_steps)
+            cosine_steps = self.total_training_steps - warmup_steps
+            
+            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 
+                start_factor=0.1,
+                end_factor=1.0,
+                total_iters=warmup_steps
+            )
+            cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, 
+                T_max=cosine_steps, 
+                eta_min=self.final_lr
+            )
+            scheduler = torch.optim.lr_scheduler.SequentialLR(
+                optimizer,
+                schedulers=[warmup_scheduler, cosine_scheduler],
+                milestones=[warmup_steps]
+            )
+            return [optimizer], [{"scheduler": scheduler, "interval": "step"}]
+        else:
+            return [optimizer]
+
+
+
+class SSLVitalSignsRegressor(SSLFineTuner):
+    """SSL Finetuner for vital signs regression (HR, SPO2). Uses MSE loss."""
+    def __init__(self,
+        backbones,
+        use_which_backbone,
+        config = None,
+        in_features: int = 256,
+        num_classes: int = 1,
+        target_names: list = None,
+        dropout: float = 0.0,
+        **kwargs
+    ) -> None:
+        kwargs['loss_type'] = 'ce'
+
+        super().__init__(
+            backbones=backbones,
+            use_which_backbone=use_which_backbone,
+            config=config,
+            in_features=in_features,
+            num_classes=2,
+            dropout=dropout,
+            **kwargs
+        )
+
+        self.num_targets = num_classes
+        self.target_names = target_names or [f"target_{i}" for i in range(num_classes)]
+        self.criterion = nn.MSELoss()
+
+        in_feat = self.linear_layer[1].in_features
+        self.linear_layer = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(in_feat, num_classes)
+        )
+
+        del self.train_pred_metrics, self.val_pred_metrics, self.test_pred_metrics
+        del self.train_prob_metrics, self.val_prob_metrics, self.test_prob_metrics
+        del self.train_pred_metrics_c, self.val_pred_metrics_c, self.test_pred_metrics_c
+        del self.train_prob_metrics_c, self.val_prob_metrics_c, self.test_prob_metrics_c
+
+    def shared_step(self, batch):
+        """Override: regression loss instead of classification."""
+        context = torch.no_grad() if not self.finetune_backbone else torch.enable_grad()
+        
+        with context:
+            psg = batch['psg']
+            if self.use_which_backbone == 'ecg':
+                x = psg[:, 0:1, :]
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'resp':
+                x = psg[:, 1:5, :]
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'elect':
+                x = psg[:, 5:, :]
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'all':
+                x = psg
+                feats = self._get_features(self.backbone, x)
+            elif self.use_which_backbone == 'fusion':
+                feats_list = []
+                if 'ecg' in self.backbones:
+                    f_ecg = self._get_features(self.backbones['ecg'], psg[:, 0:1, :])
+                    feats_list.append(f_ecg)
+                if 'resp' in self.backbones:
+                    f_resp = self._get_features(self.backbones['resp'], psg[:, 1:5, :])
+                    feats_list.append(f_resp)
+                if 'elect' in self.backbones:
+                    f_elect = self._get_features(self.backbones['elect'], psg[:, 5:, :])
+                    feats_list.append(f_elect)
+                
+                feats = torch.stack(feats_list, dim=0).mean(dim=0)
+            else:
+                raise ValueError(f"Unknown use_which_backbone: {self.use_which_backbone}")
+        
+        y = batch["label"].float()  # [B, num_targets]
+        feats = feats.view(feats.size(0), -1)
+        preds = self.linear_layer(feats)  # [B, num_targets]
+        
+        loss = self.criterion(preds, y)
+        return loss, preds, y
+
+    def training_step(self, batch, batch_idx):
+        """Override: regression metrics."""
+        loss, preds, y = self.shared_step(batch)
+        
+        with torch.no_grad():
+            for i, name in enumerate(self.target_names):
+                mae = F.l1_loss(preds[:, i], y[:, i])
+                self.log(f"train_{name}_mae", mae, on_step=False, on_epoch=True, sync_dist=True)
+        
+        self.log("train_loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True)
+        return loss
+
+    def on_train_epoch_end(self):
+        """Override: no classification metrics to compute."""
+        pass
+
+    def validation_step(self, batch, batch_idx):
+        """Override: regression metrics."""
+        loss, preds, y = self.shared_step(batch)
+        
+        for i, name in enumerate(self.target_names):
+            mae = F.l1_loss(preds[:, i], y[:, i])
+            self.log(f"val_{name}_mae", mae, on_step=False, on_epoch=True, sync_dist=True)
+        
+        overall_mae = F.l1_loss(preds, y)
+        self.log("val_loss", loss, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("val_mae", overall_mae, prog_bar=True, on_step=False, on_epoch=True, sync_dist=True)
+        return loss
+
+    def on_validation_epoch_end(self):
+        """Override: no classification metrics to compute."""
+        pass
+
+    def test_step(self, batch, batch_idx):
+        """Override: regression metrics."""
+        loss, preds, y = self.shared_step(batch)
+        
+        for i, name in enumerate(self.target_names):
+            p, t = preds[:, i], y[:, i]
+            mae = F.l1_loss(p, t)
+            mse = F.mse_loss(p, t)
+            rmse = torch.sqrt(mse)
+            
+            self.log(f"test_{name}_mae", mae, on_step=False, on_epoch=True, sync_dist=True)
+            self.log(f"test_{name}_mse", mse, on_step=False, on_epoch=True, sync_dist=True)
+            self.log(f"test_{name}_rmse", rmse, on_step=False, on_epoch=True, sync_dist=True)
+        
+        overall_mae = F.l1_loss(preds, y)
+        overall_mse = F.mse_loss(preds, y)
+        self.log("test_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_mae", overall_mae, on_step=False, on_epoch=True, sync_dist=True)
+        self.log("test_mse", overall_mse, on_step=False, on_epoch=True, sync_dist=True)
+        return loss
+
+    def on_test_epoch_end(self):
+        """Override: no classification metrics to compute."""
+        pass
+
+
+class SupervisedVitalSignsRegressor(SSLVitalSignsRegressor):
+    """Supervised from-scratch regression. Equivalent to SSLVitalSignsRegressor with finetune_backbone=True."""
+    def __init__(self, 
+        backbones,
+        use_which_backbone,
+        epochs: int = 100,
+        **kwargs
+    ):
+        kwargs['finetune_backbone'] = True
+        super().__init__(
+            backbones=backbones,
+            use_which_backbone=use_which_backbone,
+            epochs=epochs,
+            **kwargs
+        )

osf/utils/openclip_loss.py

@@ -0,0 +1,472 @@
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def get_clip_metrics(image_features, text_features, logit_scale):
+    metrics = {}
+    logits_per_image = (logit_scale * image_features @ text_features.t()).detach().cpu()
+    logits_per_text = logits_per_image.t().detach().cpu()
+
+    logits = {"image_to_text": logits_per_image, "text_to_image": logits_per_text}
+    ground_truth = torch.arange(len(text_features)).view(-1, 1)
+
+    for name, logit in logits.items():
+        ranking = torch.argsort(logit, descending=True)
+        preds = torch.where(ranking == ground_truth)[1]
+        preds = preds.detach().cpu().numpy()
+        metrics[f"{name}_mean_rank"] = preds.mean() + 1
+        metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
+        for k in [1, 5, 10]:
+            metrics[f"{name}_R@{k}"] = np.mean(preds < k)
+
+    return metrics
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        if gather_with_grad:
+            all_image_features = hvd.allgather(image_features)
+            all_text_features = hvd.allgather(text_features)
+        else:
+            with torch.no_grad():
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+    else:
+        if gather_with_grad:
+            all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+        else:
+            gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+            dist.all_gather(gathered_image_features, image_features)
+            dist.all_gather(gathered_text_features, text_features)
+            if not local_loss:
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+            all_image_features = torch.cat(gathered_image_features, dim=0)
+            all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=True,
+            gather_with_grad=True,
+            cache_labels=True,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.use_horovod = use_horovod
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, num_logits) -> torch.Tensor:
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale, return_gather_features=False):
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features,
+                text_features,
+                local_loss=self.local_loss,
+                gather_with_grad=self.gather_with_grad,
+                rank=self.rank,
+                world_size=self.world_size,
+                use_horovod=self.use_horovod,
+            )
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+            
+            if return_gather_features:
+                return logits_per_image, logits_per_text, all_image_features, all_text_features
+            else:
+                return logits_per_image, logits_per_text
+
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+        
+            return logits_per_image, logits_per_text
+
+    def forward(self, image_features, text_features, logit_scale, output_dict=False):
+        device = image_features.device
+        logits_per_image, logits_per_text = self.get_logits(image_features, text_features, logit_scale)
+
+        labels = self.get_ground_truth(device, logits_per_image.shape[0])
+
+        total_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        return {"contrastive_loss": total_loss} if output_dict else total_loss
+
+
+class CoCaLoss(ClipLoss):
+    def __init__(
+            self,
+            caption_loss_weight,
+            clip_loss_weight,
+            pad_id=0,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+    ):
+        super().__init__(
+            local_loss=local_loss,
+            gather_with_grad=gather_with_grad,
+            cache_labels=cache_labels,
+            rank=rank,
+            world_size=world_size,
+            use_horovod=use_horovod
+        )
+
+        self.clip_loss_weight = clip_loss_weight
+        self.caption_loss_weight = caption_loss_weight
+        self.caption_loss = nn.CrossEntropyLoss(ignore_index=pad_id)
+
+    def forward(self, image_features, text_features, logits, labels, logit_scale, output_dict=False):
+        if self.clip_loss_weight:
+            clip_loss = super().forward(image_features, text_features, logit_scale)
+            clip_loss = self.clip_loss_weight * clip_loss
+        else:
+            clip_loss = torch.tensor(0, device=logits.device)
+
+        caption_loss = self.caption_loss(
+            logits.permute(0, 2, 1),
+            labels,
+        )
+        caption_loss = caption_loss * self.caption_loss_weight
+
+        if output_dict:
+            return {"contrastive_loss": clip_loss, "caption_loss": caption_loss}
+
+        return clip_loss, caption_loss
+
+
+class DistillClipLoss(ClipLoss):
+
+    def dist_loss(self, teacher_logits, student_logits):
+        return -(teacher_logits.softmax(dim=1) * student_logits.log_softmax(dim=1)).sum(dim=1).mean(dim=0)
+
+    def forward(
+            self,
+            image_features,
+            text_features,
+            logit_scale,
+            dist_image_features,
+            dist_text_features,
+            dist_logit_scale,
+            output_dict=False,
+    ):
+        logits_per_image, logits_per_text = \
+            self.get_logits(image_features, text_features, logit_scale)
+
+        dist_logits_per_image, dist_logits_per_text = \
+            self.get_logits(dist_image_features, dist_text_features, dist_logit_scale)
+
+        labels = self.get_ground_truth(image_features.device, logits_per_image.shape[0])
+
+        contrastive_loss = (
+            F.cross_entropy(logits_per_image, labels) +
+            F.cross_entropy(logits_per_text, labels)
+        ) / 2
+
+        distill_loss = (
+            self.dist_loss(dist_logits_per_image, logits_per_image) +
+            self.dist_loss(dist_logits_per_text, logits_per_text)
+        ) / 2
+
+        if output_dict:
+            return {"contrastive_loss": contrastive_loss, "distill_loss": distill_loss}
+
+        return contrastive_loss, distill_loss
+
+
+def neighbour_exchange(from_rank, to_rank, tensor, group=None):
+    tensor_recv = torch.zeros_like(tensor)
+    send_op = torch.distributed.P2POp(
+        torch.distributed.isend,
+        tensor,
+        to_rank,
+        group=group,
+    )
+    recv_op = torch.distributed.P2POp(
+        torch.distributed.irecv,
+        tensor_recv,
+        from_rank,
+        group=group,
+    )
+    reqs = torch.distributed.batch_isend_irecv([send_op, recv_op])
+    for req in reqs:
+        req.wait()
+    return tensor_recv
+
+
+def neighbour_exchange_bidir(left_rank, right_rank, tensor_to_left, tensor_to_right, group=None):
+    tensor_from_left = torch.zeros_like(tensor_to_right)
+    tensor_from_right = torch.zeros_like(tensor_to_left)
+    send_op_left = torch.distributed.P2POp(
+        torch.distributed.isend,
+        tensor_to_left,
+        left_rank,
+        group=group,
+    )
+    send_op_right = torch.distributed.P2POp(
+        torch.distributed.isend,
+        tensor_to_right,
+        right_rank,
+        group=group,
+    )
+    recv_op_left = torch.distributed.P2POp(
+        torch.distributed.irecv,
+        tensor_from_left,
+        left_rank,
+        group=group,
+    )
+    recv_op_right = torch.distributed.P2POp(
+        torch.distributed.irecv,
+        tensor_from_right,
+        right_rank,
+        group=group,
+    )
+    reqs = torch.distributed.batch_isend_irecv([send_op_right, send_op_left, recv_op_right, recv_op_left])
+    for req in reqs:
+        req.wait()
+    return tensor_from_right, tensor_from_left
+
+
+class NeighbourExchange(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, from_rank, to_rank, group, tensor):
+        ctx.group = group
+        ctx.from_rank = from_rank
+        ctx.to_rank = to_rank
+        return neighbour_exchange(from_rank, to_rank, tensor, group=group)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return (None, None, None) + (NeighbourExchange.apply(ctx.to_rank, ctx.from_rank, ctx.group, grad_output),)
+
+
+def neighbour_exchange_with_grad(from_rank, to_rank, tensor, group=None):
+    return NeighbourExchange.apply(from_rank, to_rank, group, tensor)
+
+
+class NeighbourExchangeBidir(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, left_rank, right_rank, group, tensor_to_left, tensor_to_right):
+        ctx.group = group
+        ctx.left_rank = left_rank
+        ctx.right_rank = right_rank
+        return neighbour_exchange_bidir(left_rank, right_rank, tensor_to_left, tensor_to_right, group=group)
+
+    @staticmethod
+    def backward(ctx, *grad_outputs):
+        return (None, None, None) + \
+            NeighbourExchangeBidir.apply(ctx.right_rank, ctx.left_rank, ctx.group, *grad_outputs)
+
+
+def neighbour_exchange_bidir_with_grad(left_rank, right_rank, tensor_to_left, tensor_to_right, group=None):
+    return NeighbourExchangeBidir.apply(left_rank, right_rank, group, tensor_to_left, tensor_to_right)
+
+
+class SigLipLoss(nn.Module):
+    """ Sigmoid Loss for Language Image Pre-Training (SigLIP) - https://arxiv.org/abs/2303.15343
+
+    @article{zhai2023sigmoid,
+      title={Sigmoid loss for language image pre-training},
+      author={Zhai, Xiaohua and Mustafa, Basil and Kolesnikov, Alexander and Beyer, Lucas},
+      journal={arXiv preprint arXiv:2303.15343},
+      year={2023}
+    }
+    """
+    def __init__(
+            self,
+            cache_labels: bool = False,
+            rank: int = 0,
+            world_size: int = 1,
+            dist_impl: Optional[str] = None,
+    ):
+        super().__init__()
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.dist_impl = dist_impl or 'bidir'  # default to bidir exchange for now, this will likely change
+        assert self.dist_impl in ('bidir', 'shift', 'reduce', 'gather')
+
+        # FIXME: cache not currently used
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def get_ground_truth(self, device, dtype, num_logits, negative_only=False) -> torch.Tensor:
+        labels = -torch.ones((num_logits, num_logits), device=device, dtype=dtype)
+        if not negative_only:
+            labels = 2 * torch.eye(num_logits, device=device, dtype=dtype) + labels
+        return labels
+
+    def get_logits(self, image_features, text_features, logit_scale, logit_bias=None):
+        logits = logit_scale * image_features @ text_features.T
+        if logit_bias is not None:
+            logits += logit_bias
+        return logits
+
+    def _loss(self, image_features, text_features, logit_scale, logit_bias=None, negative_only=False):
+        logits = self.get_logits(image_features, text_features, logit_scale, logit_bias)
+        labels = self.get_ground_truth(
+            image_features.device,
+            image_features.dtype,
+            image_features.shape[0],
+            negative_only=negative_only,
+        )
+        loss = -F.logsigmoid(labels * logits).sum() / image_features.shape[0]
+        return loss
+
+    def forward(self, image_features, text_features, logit_scale, logit_bias, output_dict=False):
+        loss = self._loss(image_features, text_features, logit_scale, logit_bias)
+
+        if self.world_size > 1:
+            if self.dist_impl == 'bidir':
+                right_rank = (self.rank + 1) % self.world_size
+                left_rank = (self.rank - 1 + self.world_size) % self.world_size
+                text_features_to_right = text_features_to_left = text_features
+                num_bidir, remainder = divmod(self.world_size - 1, 2)
+                for i in range(num_bidir):
+                    text_features_recv = neighbour_exchange_bidir_with_grad(
+                        left_rank,
+                        right_rank,
+                        text_features_to_left,
+                        text_features_to_right,
+                    )
+                    for f in text_features_recv:
+                        loss += self._loss(
+                            image_features,
+                            f,
+                            logit_scale,
+                            logit_bias,
+                            negative_only=True,
+                        )
+                    text_features_to_left, text_features_to_right = text_features_recv
+
+                if remainder:
+                    text_features_recv = neighbour_exchange_with_grad(
+                        left_rank,
+                        right_rank,
+                        text_features_to_right
+                    )
+                    loss += self._loss(
+                        image_features,
+                        text_features_recv,
+                        logit_scale,
+                        logit_bias,
+                        negative_only=True,
+                    )
+            elif self.dist_impl == "shift":
+                right_rank = (self.rank + 1) % self.world_size
+                left_rank = (self.rank - 1 + self.world_size) % self.world_size
+                text_features_to_right = text_features
+                for i in range(self.world_size - 1):
+                    text_features_from_left = neighbour_exchange_with_grad(
+                        left_rank,
+                        right_rank,
+                        text_features_to_right,
+                    )
+                    loss += self._loss(
+                        image_features,
+                        text_features_from_left,
+                        logit_scale,
+                        logit_bias,
+                        negative_only=True,
+                    )
+                    text_features_to_right = text_features_from_left
+            elif self.dist_impl == "reduce":
+                for i in range(self.world_size):
+                    text_from_other = torch.distributed.nn.all_reduce(
+                        text_features * (self.rank == i),
+                        torch.distributed.ReduceOp.SUM,
+                    )
+                    loss += float(i != self.rank) * self._loss(
+                        image_features,
+                        text_from_other,
+                        logit_scale,
+                        logit_bias,
+                        negative_only=True,
+                    )
+            elif self.dist_impl == "gather":
+                all_text = torch.distributed.nn.all_gather(text_features)
+                for i in range(self.world_size):
+                    loss += float(i != self.rank) * self._loss(
+                        image_features,
+                        all_text[i],
+                        logit_scale,
+                        logit_bias,
+                        negative_only=True,
+                    )
+            else:
+                assert False
+
+        return {"contrastive_loss": loss} if output_dict else loss

osf/utils/results_utils.py

@@ -0,0 +1,289 @@
+"""
+Utility functions for saving experiment results to JSON/CSV.
+"""
+
+import os
+import json
+import glob
+import numpy as np
+import pandas as pd
+from typing import Dict, Any, Optional, List
+
+
+def convert_to_serializable(value):
+    """Convert tensor/numpy values to Python native types for JSON serialization."""
+    if hasattr(value, 'item'):  # torch.Tensor
+        return float(value.item())
+    elif isinstance(value, (np.ndarray, np.generic)):
+        return float(value)
+    return value
+
+
+def extract_embedding_type(root_dir: str) -> str:
+    """
+    Extract embedding type identifier from root_dir path.
+    
+    Examples:
+        ".../dino_stage1_emb_no_norm" -> "dino_no_norm"
+        ".../dino_stage1_emb" -> "dino"
+        ".../mae_emb_normalized" -> "mae_normalized"
+    
+    Args:
+        root_dir: Path to embedding directory
+        
+    Returns:
+        Short embedding type identifier
+    """
+    if not root_dir:
+        return "unknown"
+    
+    basename = os.path.basename(root_dir.rstrip('/'))
+    
+    # Remove common suffixes/patterns
+    emb_type = basename
+    emb_type = emb_type.replace("_stage1_emb", "")
+    emb_type = emb_type.replace("_stage1", "")
+    emb_type = emb_type.replace("_emb", "")
+    emb_type = emb_type.replace("final_", "")
+    
+    # Keep it concise
+    if len(emb_type) > 30:
+        emb_type = emb_type[:30]
+    
+    return emb_type if emb_type else "emb"
+
+
+def format_lr(lr: float) -> str:
+    """Format learning rate for filenames (e.g., 0.001 -> 1e-3)."""
+    if lr >= 1:
+        return f"{lr:.0f}"
+    elif lr >= 0.1:
+        return f"{lr:.1f}"
+    else:
+        # Convert to scientific notation
+        exp = 0
+        val = lr
+        while val < 1:
+            val *= 10
+            exp += 1
+        return f"{val:.0f}e-{exp}"
+
+
+def save_results_to_json(
+    test_metrics: Dict[str, Any],
+    hparams: Any,
+    extension: str,
+    ckpt_dir: str,
+    timestamp: str,
+    results_dir: str = "./results",
+    extra_fields: Optional[Dict[str, Any]] = None,
+    filename_prefix: str = "",
+) -> str:
+    """
+    Save test results to a JSON file.
+    
+    Args:
+        test_metrics: Dictionary of test metrics from trainer.test()
+        hparams: Hyperparameters namespace/object
+        extension: Experiment extension string (now used as run_name)
+        ckpt_dir: Checkpoint directory path
+        timestamp: Timestamp string
+        results_dir: Directory to save results (default: ./results)
+        extra_fields: Additional fields to include in the result record
+            - Should include: exp_type, task, dataset, model, etc.
+        filename_prefix: Prefix for the filename
+    
+    Returns:
+        Path to the saved JSON file
+    """
+    os.makedirs(results_dir, exist_ok=True)
+    
+    # Base result record
+    result_record = {
+        "run_name": extension,  # Renamed from "extension" for clarity
+        "ckpt_dir": ckpt_dir,
+        "timestamp": timestamp,
+    }
+    
+    common_fields = [
+        "model_name", "downstream_dataset_name", "ckpt_path", "stage2_ckpt_path",
+        "eval_label", "patient_cols", "use_which_backbone", "variant",
+        "in_features", "train_data_pct", "lr", "batch_size", 
+        "max_epochs", "max_steps", "loss_type", "use_mean_pool",
+        "root_dir", "is_pretrain", "pooling", "use_transformer", "use_mil",
+        "encoder_name", "encoder", "mask_channels", "encoder_size",
+        "num_classes", "seed",
+    ]
+    for field in common_fields:
+        if hasattr(hparams, field):
+            result_record[field] = getattr(hparams, field)
+    
+    standard_metrics = [
+        "test_acc", "test_f1", "test_f1_w", "test_auc", "test_auprc",
+        "test_kappa", "test_rec_m", "test_loss",
+        "test/acc", "test/f1_macro", "test/auc_macro", "test/auprc_macro",
+    ]
+    for metric in standard_metrics:
+        if metric in test_metrics:
+            key = metric.replace("/", "_")
+            result_record[key] = test_metrics[metric]
+    
+    for key, value in test_metrics.items():
+        if key.startswith("test/") or key.startswith("test_"):
+            normalized_key = key.replace("/", "_")
+            if normalized_key not in result_record:
+                result_record[normalized_key] = value
+    
+    if extra_fields:
+        result_record.update(extra_fields)
+    
+    for key, value in result_record.items():
+        result_record[key] = convert_to_serializable(value)
+    
+    if filename_prefix:
+        result_filename = f"{filename_prefix}_{timestamp}.json"
+    else:
+        model_name = getattr(hparams, 'model_name', 'model')
+        dataset_name = getattr(hparams, 'downstream_dataset_name', 'dataset')
+        label = getattr(hparams, 'eval_label', None) or getattr(hparams, 'patient_cols', 'task')
+        result_filename = f"{model_name}_{dataset_name}_{label}_{timestamp}.json"
+    
+    result_path = os.path.join(results_dir, result_filename)
+    
+    # Save to JSON
+    with open(result_path, 'w') as f:
+        json.dump(result_record, f, indent=2)
+    
+    print(f"\n{'='*80}")
+    print(f"Results saved to: {result_path}")
+    print(f"{'='*80}\n")
+    
+    return result_path
+
+
+def aggregate_results_to_csv(
+    results_dirs: List[str],
+    output_path: str = "./results/aggregated_results.csv",
+    key_columns: Optional[List[str]] = None,
+    metric_columns: Optional[List[str]] = None,
+) -> pd.DataFrame:
+    """
+    Aggregate all JSON result files from multiple directories into a single CSV.
+    
+    Args:
+        results_dirs: List of directories containing JSON result files
+        output_path: Path to save the aggregated CSV
+        key_columns: Columns to use as identifiers (default: common experiment params)
+        metric_columns: Metric columns to include (default: all test metrics)
+    
+    Returns:
+        DataFrame with aggregated results
+    """
+    if key_columns is None:
+        key_columns = [
+            "exp_type", "task", "dataset", "model", "encoder",
+            "train_data_pct", "lr", "embedding_type",
+            "pretrain_ckpt_path", "finetuned_ckpt_dir", "trained_ckpt_dir",
+            "stage2_pretrain_ckpt", "embedding_root_dir",
+            "model_name", "downstream_dataset_name", "eval_label", "patient_cols",
+            "use_which_backbone", "variant", "loss_type",
+            "use_mean_pool", "pooling", "use_transformer", "use_mil",
+            "mask_channels", "mask_channels_str",
+            "ckpt_path", "stage2_ckpt_path", "root_dir",
+        ]
+    
+    if metric_columns is None:
+        metric_columns = [
+            "test_acc", "test_f1", "test_f1_w", "test_auc", "test_auprc",
+            "test_kappa", "test_rec_m", "test_loss",
+        ]
+    
+    all_records = []
+    
+    for results_dir in results_dirs:
+        if not os.path.exists(results_dir):
+            print(f"[WARN] Directory not found: {results_dir}")
+            continue
+        
+        # Find all JSON files
+        json_files = glob.glob(os.path.join(results_dir, "*.json"))
+        print(f"[INFO] Found {len(json_files)} JSON files in {results_dir}")
+        
+        for json_file in json_files:
+            try:
+                with open(json_file, 'r') as f:
+                    record = json.load(f)
+                record['_source_file'] = os.path.basename(json_file)
+                record['_source_dir'] = results_dir
+                all_records.append(record)
+            except Exception as e:
+                print(f"[WARN] Failed to load {json_file}: {e}")
+    
+    if not all_records:
+        print("[WARN] No records found!")
+        return pd.DataFrame()
+    
+    # Convert to DataFrame
+    df = pd.DataFrame(all_records)
+    
+    existing_key_cols = [c for c in key_columns if c in df.columns]
+    existing_metric_cols = [c for c in metric_columns if c in df.columns]
+    
+    per_class_cols = [c for c in df.columns if c.startswith("test_") and c not in existing_metric_cols]
+    per_class_cols = sorted(per_class_cols)
+    
+    other_cols = [c for c in df.columns if c not in existing_key_cols + existing_metric_cols + per_class_cols]
+    
+    ordered_cols = existing_key_cols + existing_metric_cols + per_class_cols + other_cols
+    df = df[[c for c in ordered_cols if c in df.columns]]
+    
+    os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
+    df.to_csv(output_path, index=False)
+    
+    print(f"\n{'='*80}")
+    print(f"Aggregated {len(all_records)} results to: {output_path}")
+    print(f"Columns: {list(df.columns[:10])}... ({len(df.columns)} total)")
+    print(f"{'='*80}\n")
+    
+    return df
+
+
+def load_results_from_json(json_path: str) -> Dict[str, Any]:
+    """Load a single JSON result file."""
+    with open(json_path, 'r') as f:
+        return json.load(f)
+
+
+def filter_results(
+    df: pd.DataFrame,
+    model_name: Optional[str] = None,
+    dataset_name: Optional[str] = None,
+    eval_label: Optional[str] = None,
+    patient_cols: Optional[str] = None,
+) -> pd.DataFrame:
+    """
+    Filter aggregated results DataFrame by common fields.
+    
+    Args:
+        df: DataFrame from aggregate_results_to_csv()
+        model_name: Filter by model name
+        dataset_name: Filter by downstream dataset name
+        eval_label: Filter by eval label (stage 1)
+        patient_cols: Filter by patient columns (stage 2)
+    
+    Returns:
+        Filtered DataFrame
+    """
+    filtered = df.copy()
+    
+    if model_name is not None and 'model_name' in filtered.columns:
+        filtered = filtered[filtered['model_name'] == model_name]
+    if dataset_name is not None and 'downstream_dataset_name' in filtered.columns:
+        filtered = filtered[filtered['downstream_dataset_name'] == dataset_name]
+    if eval_label is not None and 'eval_label' in filtered.columns:
+        filtered = filtered[filtered['eval_label'] == eval_label]
+    if patient_cols is not None and 'patient_cols' in filtered.columns:
+        filtered = filtered[filtered['patient_cols'] == patient_cols]
+    
+    return filtered
+

osf_backbone.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c51190b1942556969af3c3d63c2e59430ddb1ea0377c50ea87df83712fc31857
+size 341360652

pretrained_weights/readme.md

@@ -0,0 +1 @@
+Please download the checkpoint throught the link in the readme of the directory root.

requirements.txt

@@ -0,0 +1,193 @@
+absl-py==2.3.1
+accelerate==1.2.1
+aiohappyeyeballs==2.4.4
+aiohttp==3.11.10
+aiosignal==1.3.1
+albucore==0.0.24
+albumentations==2.0.8
+altair==5.5.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+asttokens==3.0.0
+async-timeout==5.0.1
+attrs==25.3.0
+beartype==0.22.2
+bitarray==3.0.0
+blinker==1.9.0
+braceexpand==0.1.7
+certifi==2025.10.5
+cffi==1.17.1
+charset-normalizer==3.4.3
+click==8.1.7
+coloredlogs==15.0.1
+comm==0.2.3
+contourpy==1.3.1
+cosine_annealing_warmup @ git+https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup@12d03c07553aedd3d9e9155e2b3e31ce8c64081a
+cycler==0.12.1
+Cython==3.0.11
+datasets==3.2.0
+debugpy==1.8.17
+decorator==5.2.1
+diffusers==0.32.1
+dill==0.3.8
+docker-pycreds==0.4.0
+easydict==1.13
+efficientnet_pytorch==0.7.1
+einops==0.8.0
+ema-pytorch==0.7.7
+et_xmlfile==2.0.0
+exceptiongroup==1.3.0
+executing==2.2.1
+fairseq_signals_backbone @ git+https://github.com/fuying-wang/fairseq-signals@27d94bab8a1040879c011609df1488aac21a586a
+filelock==3.20.0
+flatbuffers==25.9.23
+fonttools==4.55.1
+frozenlist==1.5.0
+fsspec==2024.6.1
+gitdb==4.0.11
+GitPython==3.1.43
+grpcio==1.75.1
+h5py==3.14.0
+hf-xet==1.1.10
+huggingface-hub==0.34.4
+humanfriendly==10.0
+hydra-core==1.3.2
+idna==3.10
+imageio==2.37.0
+importlib_metadata==8.7.0
+insightface==0.7.3
+ipdb==0.13.13
+ipykernel==6.30.1
+ipython==8.37.0
+jedi==0.19.2
+Jinja2==3.1.6
+joblib==1.4.2
+jupyter_client==8.6.3
+jupyter_core==5.8.1
+kiwisolver==1.4.7
+kornia==0.8.1
+kornia_rs==0.1.9
+lazy_loader==0.4
+lightning==2.4.0
+lightning-utilities==0.11.9
+llvmlite==0.46.0
+loguru==0.7.3
+lxml==5.3.0
+Markdown==3.9
+MarkupSafe==3.0.3
+matplotlib==3.9.3
+matplotlib-inline==0.1.7
+ml_dtypes==0.5.3
+mne==1.10.1
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+munch==4.0.0
+narwhals==2.6.0
+nest-asyncio==1.6.0
+networkx==3.4.2
+neurokit2==0.2.12
+ninja==1.13.0
+nltk==3.9.1
+numba==0.63.1
+numpy==2.1.2
+omegaconf==2.3.0
+onnx==1.19.1
+onnx2torch==1.5.15
+onnxruntime==1.23.1
+opencv-python==4.12.0.88
+opencv-python-headless==4.12.0.88
+openpyxl==3.1.5
+packaging==24.2
+pandas==2.2.3
+parso==0.8.5
+peft==0.14.0
+pexpect==4.9.0
+pillow==11.0.0
+platformdirs==4.4.0
+pooch==1.8.2
+portalocker==3.0.0
+POT==0.9.5
+pretrainedmodels==0.7.4
+prettytable==3.16.0
+prompt_toolkit==3.0.52
+propcache==0.2.1
+protobuf==5.29.1
+psutil==7.1.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pyarrow==18.1.0
+pycparser==2.23
+pydantic==2.10.3
+pydantic_core==2.27.1
+pydeck==0.9.1
+Pygments==2.19.2
+pynndescent==0.5.13
+pyparsing==3.2.0
+pysam==0.23.3
+python-dateutil==2.9.0.post0
+pytorch-lightning==2.4.0
+pytorch-warmup==0.2.0
+pytz==2024.2
+PyWavelets==1.8.0
+PyYAML==6.0.3
+pyzmq==27.1.0
+regex==2024.11.6
+requests==2.32.5
+sacrebleu==2.4.3
+safetensors==0.6.2
+scikit-image==0.25.2
+scikit-learn==1.7.2
+scipy==1.14.1
+seaborn==0.13.2
+segmentation_models_pytorch==0.4.0
+sentencepiece==0.2.1
+sentry-sdk==2.19.2
+setproctitle==1.3.4
+simsimd==6.5.3
+six==1.17.0
+smmap==5.0.1
+soundfile==0.12.1
+stack-data==0.6.3
+streamlit==1.50.0
+stringzilla==4.2.1
+sympy==1.13.1
+tabulate==0.9.0
+tenacity==9.1.2
+tensorboard==2.20.0
+tensorboard-data-server==0.7.2
+tensorboardX==2.6.4
+threadpoolctl==3.5.0
+tifffile==2025.5.10
+timm==1.0.12
+tokenizers==0.21.0
+toml==0.10.2
+tomli==2.2.1
+torch==2.5.1
+torchaudio==2.5.1
+torchdiffeq==0.2.5
+torchmetrics==1.6.0
+torchtools @ git+https://github.com/pabloppp/pytorch-tools@610158d5016d6418aee27f956e7afd17ff35ba04
+torchvision==0.20.1
+tornado==6.5.2
+tqdm==4.67.1
+traitlets==5.14.3
+transformers==4.47.0
+typing-inspection==0.4.1
+typing_extensions==4.15.0
+tzdata==2024.2
+umap-learn==0.5.9.post2
+unet==0.8.1
+urllib3==2.5.0
+vitaldb==1.5.8
+wandb==0.21.4
+warmup_scheduler==0.3
+watchdog==6.0.0
+wcwidth==0.2.13
+webdataset==1.0.2
+Werkzeug==3.1.3
+wfdb==4.1.2
+xgboost==2.1.3
+xxhash==3.5.0
+yarl==1.18.3
+zipp==3.23.0

train_config.py

@@ -0,0 +1,32 @@
+from config import *
+
+
+# Uni-encoder models (simclr, dino, mae, vqvae, ar, etc.)
+TRAIN_EDF_COLS_UNI_ENC = [ECG, EMG_Chin, EMG_LLeg, EMG_RLeg,
+                  ABD, THX, NP, SN, 
+                  EOG_E1_A2, EOG_E2_A1,EEG_C3_A2, EEG_C4_A1,
+            ]
+TRAIN_EDF_COLS_MULTI_ENC = [ECG, 
+                  ABD, THX, NP, SN, 
+                  EMG_Chin, EMG_LLeg, EMG_RLeg,
+                  EOG_E1_A2, EOG_E2_A1,EEG_C3_A2, EEG_C4_A1,
+            ]
+TRAIN_EDF_COLS_TYPE3 = [ECG, ABD, THX, NP, SN]
+TRAIN_EDF_COLS_TYPE4 = [ECG, ABD, THX]
+
+
+MONITOR_TYPE_MAP = {
+    "main": TRAIN_EDF_COLS_UNI_ENC,
+    "type3": TRAIN_EDF_COLS_TYPE3,
+    "type4": TRAIN_EDF_COLS_TYPE4,
+}
+STAGE2_LABEL_PATH_WITH_PATHHEAD = "/path/to/your/label/splits"
+CKPT_PATH = "/path/to/your/checkpoints"
+MODEL_LIST = ["dino_ours"]
+
+AUGMENTATION_MAP = {
+    "dino_ours": "chan_then_pcspan",
+}
+SPLIT_DATA_FOLDER = "/path/to/your/postprocessed/data"
+PRETRAIN_VAL_DATASET_LIST = ['shhs']
+NEED_NORM_COL = [HR, SPO2, OX]