diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 74592ef..abecdb1 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -10,15 +10,17 @@ BioFoundation is both a research repository and an onboarding codebase. Contribu - Preserve the full Apache 2.0 header in every Python source and Hydra YAML file. - Normalize task input through `biofoundation.core.batch.as_signal_batch`. - Keep model-family metadata in `biofoundation/model_registry.py`, including its paper and Hugging Face repository. +- Change `biofoundation/` additively. New fields, functions, and classes are fine; existing names must not change signature or behaviour, and new defaults must leave existing callers observing no difference. A change that cannot be made additively needs a deprecation period and a major version bump of `biofoundation.__version__`. +- Record decisions that are expensive to reverse as an ADR under [`docs/adr`](docs/adr/). - Use Hydra configs for reproducible settings instead of embedding experiment paths or hyperparameters in Python. - Mark local values that a new user must supply with `#CHANGEME` and document the expected value. ## Adding a Model -1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`. +1. Add the `nn.Module` implementation under `models/` and its Hydra model config under `config/model/`. A family may either bundle its output layer into the model, as the five original families do, or separate an encoder from a prediction head under `models/model_heads/` with a matching `config/model_head/` entry. The second shape is described by the protocols in `biofoundation/core/protocols.py` and lets one pre-trained encoder serve every downstream task. 2. Add pre-training and fine-tuning experiments under `config/experiment/`. 3. Use an existing task when its behavior fits. New task steps must use the shared batch adapter. -4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, paper, and Hugging Face URL. +4. Register the model in `biofoundation/model_registry.py`, including modalities, architecture, experiment names, batch requirements, venue, paper, and Hugging Face URL. Use the case-folded display name as the key, and list any prediction heads in `head_targets`. Setting `model_family` in the experiment lets a task look its family up and enforce the declared batch requirements at runtime. Only the S-CEReBrO tasks read it today; the five bundled families validate nothing at runtime. 5. Add a model page under `docs/model/` with input assumptions, training details, and checkpoint usage. 6. Extend the contract tests for any new shared behavior. diff --git a/README.md b/README.md index e6e6de6..6c04f59 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ Authors: Thorir Mar Ingolfsson, Anna Tegon, Berkay Döner, Xiaying Wang, Matteo > **TL;DR:** Choose a model from the table below, install the training dependencies, set `DATA_PATH` and `CHECKPOINT_DIR`, then run `python -u run_train.py +experiment=_pretrain` or the matching fine-tuning experiment. Each model page links its Hugging Face weights and exact checkpoint command. ARES is separate and only needed for embedded deployment. -BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind five model families. +BioFoundation is a research and onboarding codebase for foundation models across EEG, sEMG, ECG, and PPG. It collects the model implementations, Hydra experiments, preprocessing tools, and pretrained releases behind six model families. The training stack is built on PyTorch Lightning and Hydra. Embedded deployment through ARES is maintained as a separate toolchain inside the repository. @@ -19,6 +19,7 @@ The training stack is built on PyTorch Lightning and Hydra. Embedded deployment | [TinyMyo](docs/model/TinyMyo.md) | sEMG | Rotary Transformer | [Paper](https://arxiv.org/abs/2512.15729) / [Hugging Face](https://huggingface.co/PulpBio/TinyMyo) | | [LuMamba](docs/model/LuMamba.md) | EEG | Query-unified Mamba | [Paper](https://arxiv.org/abs/2603.19100) / [Hugging Face](https://huggingface.co/PulpBio/LuMamba) | | [PanLUNA](docs/model/PanLUNA.md) | EEG, ECG, PPG | Multimodal query-unified Transformer | [Paper](https://arxiv.org/abs/2604.04297) / [Hugging Face](https://huggingface.co/PulpBio/PanLUNA) | +| [S-CEReBrO](docs/model/SCEReBrO.md) | EEG | Windowed alternating-attention Transformer | [Paper](https://arxiv.org/abs/2607.27913) / [Hugging Face](https://huggingface.co/PulpBio/S-CEReBrO) | The machine-readable [`model_registry.py`](biofoundation/model_registry.py) records the experiment names, papers, Hugging Face repositories, modalities, and batch metadata requirements for these families. @@ -74,6 +75,7 @@ Choose another `+experiment` from the model registry. Before a long run, review | [`datasets`](datasets/) | Dataset readers and sample contracts. | | [`data_module`](data_module/) | Lightning data modules and loader composition. | | [`config`](config/) | Hydra defaults, modules, and reproducible experiments. | +| [`docs/adr`](docs/adr/) | Architecture decision records for shared contracts. | | [`make_datasets`](make_datasets/) | Raw-data preprocessing and HDF5 conversion. | | [`criterion`](criterion/) | Training objectives. | | [`tests`](tests/) | Fast repository and refactoring contracts. | @@ -105,6 +107,6 @@ Numerical changes to models, losses, or datasets should also be tested with repr ## Licensing and Support -The source code is licensed under Apache 2.0. Pretrained weights in the five PulpBio Hugging Face repositories are licensed under CC BY-ND 4.0; see the model cards for terms and checkpoint-specific details. +The source code is licensed under Apache 2.0. Pretrained weights in the six PulpBio Hugging Face repositories are licensed under CC BY-ND 4.0; see the model cards for terms and checkpoint-specific details. For questions and support, open an [issue](https://github.com/pulp-bio/BioFoundation/issues). For changes, start with [CONTRIBUTING.md](CONTRIBUTING.md). diff --git a/biofoundation/__init__.py b/biofoundation/__init__.py index da9f3de..b72754c 100644 --- a/biofoundation/__init__.py +++ b/biofoundation/__init__.py @@ -17,5 +17,21 @@ #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* -"""Shared infrastructure for the BioFoundation model zoo.""" +"""Shared infrastructure for the BioFoundation model zoo. + +This package holds the contracts that model families, tasks, and datasets agree on: +the batch layout, the encoder and prediction-head protocols, checkpoint entry points, +environment validation, and the model registry. Everything here is imported by code +that must run without PyTorch installed, so this package has no heavyweight runtime +dependencies. + +Changes follow an additive rule. New fields, functions, and classes may be added; +existing names do not change signature or behaviour, and defaults are chosen so that +an existing caller observes no difference. A change that cannot be made additively is +a major version bump and needs a deprecation period first. +""" + +__version__ = "0.2.0" + +__all__ = ["__version__"] diff --git a/biofoundation/core/__init__.py b/biofoundation/core/__init__.py index 57e6448..72d7bec 100644 --- a/biofoundation/core/__init__.py +++ b/biofoundation/core/__init__.py @@ -20,10 +20,15 @@ """Stable contracts shared by model-specific implementations.""" from biofoundation.core.batch import BatchRequirements, SignalBatch, as_signal_batch, require_batch_fields +from biofoundation.core.checkpoints import SafetensorsCheckpointMixin +from biofoundation.core.protocols import PredictionHead, SignalEncoder __all__ = [ "BatchRequirements", + "PredictionHead", + "SafetensorsCheckpointMixin", "SignalBatch", + "SignalEncoder", "as_signal_batch", "require_batch_fields", ] diff --git a/biofoundation/core/batch.py b/biofoundation/core/batch.py index 5f44502..b3cc9d1 100644 --- a/biofoundation/core/batch.py +++ b/biofoundation/core/batch.py @@ -29,24 +29,55 @@ class SignalBatch(TypedDict, total=False): Only ``input`` is universally required. Model adapters can require channel or sensor metadata without forcing simpler models to manufacture it. + + Two independent electrode-geometry representations are supported, and a model + requires exactly one of them. They are peers rather than alternatives to convert + between, so a dataset declares which one it produces and the registry records + which one each model consumes: + + ``channel_locations`` + Shape ``(batch, channels, 3)``. One 3D coordinate per channel. For a bipolar + derivation this is the midpoint of the two electrodes. Consumed by LUNA, + LuMamba and PanLUNA. + + ``channel_coords`` + Shape ``(batch, channels, 2, 3)``. Both electrodes of every channel kept + separate, so a bipolar pair and a scalp-plus-reference channel stay + distinguishable. Consumed by S-CEReBrO. + + Padding metadata describes how much of a sample is filler, which lets montages of + different sizes share one batch. It is only meaningful on mapping-shaped batches; + the tuple form accepted by :func:`as_signal_batch` cannot carry it. """ input: Any label: Any channel_names: Any channel_locations: Any + channel_coords: Any sensor_type: Any + num_padded_channels: Any + num_padded_timesteps: Any metadata: Mapping[str, Any] @dataclass(frozen=True) class BatchRequirements: - """Metadata fields required by a particular model adapter.""" + """Metadata fields required by a particular model adapter. + + Every field defaults to ``False``, so a model only declares what it actually + reads. ``channel_locations`` and ``channel_coords`` are the two electrode-geometry + representations described on :class:`SignalBatch`; a model sets one of them. + """ label: bool = False channel_names: bool = False channel_locations: bool = False sensor_type: bool = False + # Appended after the original fields so positional construction keeps its meaning. + channel_coords: bool = False + num_padded_channels: bool = False + num_padded_timesteps: bool = False def as_signal_batch(batch: Any) -> SignalBatch: @@ -81,6 +112,9 @@ def require_batch_fields(batch: SignalBatch, requirements: BatchRequirements) -> ("channel_names", requirements.channel_names), ("channel_locations", requirements.channel_locations), ("sensor_type", requirements.sensor_type), + ("channel_coords", requirements.channel_coords), + ("num_padded_channels", requirements.num_padded_channels), + ("num_padded_timesteps", requirements.num_padded_timesteps), ) missing = [name for name, required in required_fields if required and name not in batch] if missing: diff --git a/biofoundation/core/checkpoints.py b/biofoundation/core/checkpoints.py new file mode 100644 index 0000000..c0c250d --- /dev/null +++ b/biofoundation/core/checkpoints.py @@ -0,0 +1,69 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* + +"""Checkpoint entry points shared by tasks that separate encoder from head. + +``run_train.py`` calls ``load_pretrained_checkpoint`` for a ``.ckpt`` and +``load_safetensors_checkpoint`` for a ``.safetensors``; this mixin maps both onto a +task's single ``load_weights``, leaving tensor matching to the task. + +``load_weights`` is not called ``load_from_checkpoint`` because that name is a +:class:`pytorch_lightning.LightningModule` classmethod, and an instance method of the +same name shadows it so ``MyTask.load_from_checkpoint(path)`` binds the path to ``self``. +""" + +from typing import Any + + +class SafetensorsCheckpointMixin: + """Expose BioFoundation's two checkpoint entry points over one ``load_weights``. + + Mix in ahead of :class:`pytorch_lightning.LightningModule`. The safetensors path + wraps the flat mapping in ``{"state_dict": ...}`` and prefixes bare keys with + ``model.``, matching how the pre-training task saves an encoder. + """ + + def load_pretrained_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any: + """Load a Lightning ``.ckpt`` by delegating to ``load_weights``.""" + + return self.load_weights(checkpoint_path=model_ckpt, **kwargs) + + def load_safetensors_checkpoint(self, model_ckpt: str, **kwargs: Any) -> Any: + """Load a ``.safetensors`` file through the same path as a Lightning checkpoint.""" + + import tempfile + from pathlib import Path + + import torch + from safetensors.torch import load_file + + state_dict = { + key if key.startswith(("model.", "model_head.")) else f"model.{key}": value + for key, value in load_file(model_ckpt).items() + } + + with tempfile.TemporaryDirectory() as directory: + converted = Path(directory) / "converted.ckpt" + torch.save({"state_dict": state_dict}, converted) + return self.load_weights(checkpoint_path=str(converted), **kwargs) + + + + +__all__ = ["SafetensorsCheckpointMixin"] diff --git a/biofoundation/core/protocols.py b/biofoundation/core/protocols.py new file mode 100644 index 0000000..bfc4ad5 --- /dev/null +++ b/biofoundation/core/protocols.py @@ -0,0 +1,83 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* + +"""Structural contracts for the encoder and prediction-head split. + +The five original families bundle their output layer into the model, selecting it from +``num_classes``. Newer families separate an encoder that produces token embeddings from +a prediction head that consumes them, each from its own Hydra group; the protocols +below describe that second shape. + +Conformance is structural, so the bundled families stay valid unchanged. The signatures +are concrete on purpose: a ``forward`` accepting ``*args, **kwargs`` would be satisfied +by every :class:`torch.nn.Module` and check nothing. + +Tensors are annotated only under :data:`typing.TYPE_CHECKING`, because the fast contract +suite imports this module without PyTorch installed. +""" + +from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable + +if TYPE_CHECKING: # pragma: no cover - typing only + from torch import Tensor +else: # pragma: no cover - runtime fallback keeps this module torch-free + Tensor = Any + + +@runtime_checkable +class SignalEncoder(Protocol): + """A model that turns a patched biosignal into contextualised token embeddings. + + The encoder owns tokenisation, positional and channel embeddings, and the + backbone. It owns no task-specific output layer, so one pre-trained encoder can + be paired with any :class:`PredictionHead` without being rebuilt. + + Implementations accept ``(batch, channels, patches, patch_size)`` and return + ``(batch, channels * patches, embed_dim)``. Keyword arguments beyond the input + carry whatever batch metadata the family requires, declared in its + :class:`~biofoundation.core.batch.BatchRequirements`. + """ + + def forward(self, x: "Tensor", channel_positions: "Tensor") -> "Tensor": + """Encode a patched biosignal into token embeddings. + + Args: + x: ``(batch, channels, patches, patch_size)``. + channel_positions: Electrode geometry, shaped as the family's + :class:`~biofoundation.core.batch.BatchRequirements` declares. + """ + ... + + +@runtime_checkable +class PredictionHead(Protocol): + """A module that turns encoder token embeddings into a task prediction. + + Heads accept ``(batch, num_tokens, embed_dim)``. The output shape is the head's + own concern: class logits, a scalar per window, or one reconstructed patch per + token. A head never reads the raw waveform and never receives batch metadata, so + swapping the task means swapping the head alone. + """ + + def forward(self, x: "Tensor") -> "Tensor": + """Map ``(batch, num_tokens, embed_dim)`` token embeddings to a prediction.""" + ... + + +__all__ = ["PredictionHead", "SignalEncoder"] diff --git a/biofoundation/model_registry.py b/biofoundation/model_registry.py index 7b801f2..86ecad8 100644 --- a/biofoundation/model_registry.py +++ b/biofoundation/model_registry.py @@ -28,7 +28,13 @@ @dataclass(frozen=True) class ModelSpec: - """Human- and machine-readable entry for a BioFoundation model family.""" + """Human- and machine-readable entry for a BioFoundation model family. + + ``model_target`` names the module a family is built from. For families that bundle + their output layer into the model, that is the whole model. For families that + separate the two, it is the encoder, and ``head_targets`` lists the prediction + heads it can be paired with. + """ display_name: str modalities: tuple[str, ...] @@ -39,6 +45,9 @@ class ModelSpec: huggingface_url: str paper_url: str batch_requirements: BatchRequirements = BatchRequirements() + venue: str = "" + head_targets: tuple[str, ...] = () + size_variants: tuple[str, ...] = () MODEL_REGISTRY: Mapping[str, ModelSpec] = MappingProxyType( @@ -52,6 +61,7 @@ class ModelSpec: finetune_experiment="FEMBA_finetune", huggingface_url="https://huggingface.co/PulpBio/FEMBA", paper_url="https://arxiv.org/abs/2502.06438", + venue="EMBC 2025", ), "luna": ModelSpec( display_name="LUNA", @@ -63,6 +73,8 @@ class ModelSpec: huggingface_url="https://huggingface.co/PulpBio/LUNA", paper_url="https://arxiv.org/abs/2510.22257", batch_requirements=BatchRequirements(channel_locations=True), + venue="NeurIPS 2025", + size_variants=("base", "large", "huge"), ), "tinymyo": ModelSpec( display_name="TinyMyo", @@ -73,6 +85,7 @@ class ModelSpec: finetune_experiment="TinyMyo_finetune", huggingface_url="https://huggingface.co/PulpBio/TinyMyo", paper_url="https://arxiv.org/abs/2512.15729", + venue="arXiv preprint", ), "lumamba": ModelSpec( display_name="LuMamba", @@ -84,6 +97,8 @@ class ModelSpec: huggingface_url="https://huggingface.co/PulpBio/LuMamba", paper_url="https://arxiv.org/abs/2603.19100", batch_requirements=BatchRequirements(channel_locations=True), + venue="EUSIPCO 2026", + size_variants=("tiny",), ), "panluna": ModelSpec( display_name="PanLUNA", @@ -98,6 +113,26 @@ class ModelSpec: channel_locations=True, sensor_type=True, ), + venue="AICAS 2026", + ), + "s-cerebro": ModelSpec( + display_name="S-CEReBrO", + modalities=("EEG",), + architecture="Windowed alternating-attention Transformer", + model_target="models.s_cerebro.SCerebroEncoder", + pretrain_experiment="SCEReBrO_pretrain", + finetune_experiment="SCEReBrO_finetune", + huggingface_url="https://huggingface.co/PulpBio/S-CEReBrO", + paper_url="https://arxiv.org/abs/2607.27913", + batch_requirements=BatchRequirements(channel_coords=True), + venue="MICCAI 2026", + head_targets=( + "models.model_heads.patch_reconstruction_head.PatchReconstructionHead", + "models.model_heads.mlp_classification_head.MlpClassificationHead", + "models.model_heads.mlp_regression_head.MlpRegressionHead", + "models.model_heads.sequence_classification_head.SequenceClassificationHead", + ), + size_variants=("tiny",), ), } ) diff --git a/config/criterion/ce_criterion.yaml b/config/criterion/ce_criterion.yaml new file mode 100644 index 0000000..0694a98 --- /dev/null +++ b/config/criterion/ce_criterion.yaml @@ -0,0 +1,23 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +criterion: + _target_: criterion.ce_criterion.CrossEntropyLossWrapper + label_smoothing: 0.0 + weight: null diff --git a/config/criterion/focal_criterion.yaml b/config/criterion/focal_criterion.yaml new file mode 100644 index 0000000..73beb5c --- /dev/null +++ b/config/criterion/focal_criterion.yaml @@ -0,0 +1,23 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +criterion: + _target_: criterion.focal_criterion.FocalLossWrapper + alpha: 0.8 + gamma: 0.7 diff --git a/config/criterion/masked_reconstruction_loss.yaml b/config/criterion/masked_reconstruction_loss.yaml new file mode 100644 index 0000000..3e431f5 --- /dev/null +++ b/config/criterion/masked_reconstruction_loss.yaml @@ -0,0 +1,23 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +criterion: + _target_: criterion.masked_reconstruction_loss.MaskedReconstructionLoss + loss_type: l2 + alpha: 0.1 # weight of the visible-patch term; 0 disables it diff --git a/config/criterion/mse_criterion.yaml b/config/criterion/mse_criterion.yaml new file mode 100644 index 0000000..25b65ac --- /dev/null +++ b/config/criterion/mse_criterion.yaml @@ -0,0 +1,21 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +criterion: + _target_: criterion.mse_criterion.MSELossWrapper diff --git a/config/data_module/finetune_data_module_SCEReBrO.yaml b/config/data_module/finetune_data_module_SCEReBrO.yaml new file mode 100644 index 0000000..a4318a7 --- /dev/null +++ b/config/data_module/finetune_data_module_SCEReBrO.yaml @@ -0,0 +1,56 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +data_module: + _target_: 'data_module.finetuning_data_module.FinetuningDataModule' + name: "eeg" + cfg: + num_workers: ${num_workers} + batch_size: ${batch_size} + + # Splits are fixed by the preprocessing step, not resampled here. Point dataset_root + # at another prepared corpus to fine-tune on it; see docs/model/SCEReBrO.md for the + # channel count, window length and class count each one expects. + train: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: ${dataset_root} + split: train + dataset_kind: ${dataset_kind} + apply_minmax: True + apply_zero_padding: False + label_mode: ${label_mode} + use_cache: False + val: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: ${dataset_root} + split: val + dataset_kind: ${dataset_kind} + apply_minmax: True + apply_zero_padding: False + label_mode: ${label_mode} + use_cache: False + test: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: ${dataset_root} + split: test + dataset_kind: ${dataset_kind} + apply_minmax: True + apply_zero_padding: False + label_mode: ${label_mode} + use_cache: False diff --git a/config/data_module/pretrain_data_module_SCEReBrO.yaml b/config/data_module/pretrain_data_module_SCEReBrO.yaml new file mode 100644 index 0000000..7c37da5 --- /dev/null +++ b/config/data_module/pretrain_data_module_SCEReBrO.yaml @@ -0,0 +1,118 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# Each corpus is read from here, as written by its make_datasets script. Relocate them +# all with pretraining_root=/some/dir. +pretraining_root: ${env:DATA_PATH}/pretraining + +data_module: + _target_: 'data_module.pretraining_data_module.PretrainingDataModule' + name: "eeg" + cfg: + num_workers: ${num_workers} + batch_size: ${batch_size} + + # Held-out fraction and the seed that fixes the split. run_train.py instantiates the + # data module without arguments, so the seed is threaded through the config rather + # than passed at call time; every rank therefore derives the same split. + val_ratio: 0.2 + seed: ${seed} + + # Corpora are concatenated. Set a corpus to null to drop it from a run. + train: + TUEG: + _target_: 'datasets.tueg_dataset.TUEGDataset' + lmdb_path: '${pretraining_root}/TUEG/TUEG.lmdb' + keys_path: '${pretraining_root}/TUEG/TUEG_keys.txt' + max_channels: ${model.max_channels} + sampling_freq: 200 + slice_duration: 30 + use_cache: False + SEED: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/SEED/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + SEED-IV: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/SEED-IV/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + SEED-GER: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/SEED-GER/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + SEED-FRA: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/SEED-FRA/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + BOAS: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/BOAS/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + GWD: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/GWD/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + SleepEDFx: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/SleepEDFx/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False + BCI-NER: + _target_: 'datasets.lmdb_dataset.LMDBDataset' + path: '${pretraining_root}/BCI-NER/' + dataset_kind: window + apply_minmax: True + apply_zero_padding: True + max_channels: ${model.max_channels} + max_timesteps: ${model.max_timesteps} + use_cache: False diff --git a/config/dataset/chb-mit.yaml b/config/dataset/chb-mit.yaml new file mode 100644 index 0000000..047a542 --- /dev/null +++ b/config/dataset/chb-mit.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# CHB-MIT: seizure detection. 16 channels, 10 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/CHB-MIT +dataset_kind: window +label_mode: classification + +model: + num_channels: 16 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 10 diff --git a/config/dataset/isruc.yaml b/config/dataset/isruc.yaml new file mode 100644 index 0000000..31cf08d --- /dev/null +++ b/config/dataset/isruc.yaml @@ -0,0 +1,38 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# ISRUC: sleep staging. 6 channels, sequences of 20 consecutive 30 s epochs, 5 classes +# per epoch. The sequence axis is folded into the batch by the task and restored by the +# head, so the encoder still sees one epoch at a time. +defaults: + - /model_head: sequence_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/ISRUC +dataset_kind: sequence +label_mode: classification + +model: + num_channels: 6 + +model_head: + num_classes: 5 + num_patches: 30 + sequence_length: 20 diff --git a/config/dataset/mental-arithmetic.yaml b/config/dataset/mental-arithmetic.yaml new file mode 100644 index 0000000..3a0fac8 --- /dev/null +++ b/config/dataset/mental-arithmetic.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# MentalArithmetic: mental arithmetic. 20 channels, 5 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/MentalArithmetic +dataset_kind: window +label_mode: classification + +model: + num_channels: 20 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 5 diff --git a/config/dataset/mumtaz.yaml b/config/dataset/mumtaz.yaml new file mode 100644 index 0000000..c6695e0 --- /dev/null +++ b/config/dataset/mumtaz.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# Mumtaz: depression. 19 channels, 5 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/Mumtaz +dataset_kind: window +label_mode: classification + +model: + num_channels: 19 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 5 diff --git a/config/dataset/neonate.yaml b/config/dataset/neonate.yaml new file mode 100644 index 0000000..eba38e4 --- /dev/null +++ b/config/dataset/neonate.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# Neonate: neonatal seizure detection. 18 channels, 5 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/Neonate +dataset_kind: window +label_mode: classification + +model: + num_channels: 18 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 5 diff --git a/config/dataset/physionet-mi.yaml b/config/dataset/physionet-mi.yaml new file mode 100644 index 0000000..31c09e2 --- /dev/null +++ b/config/dataset/physionet-mi.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# PhysioNet-MI: motor imagery. 64 channels, 4 s windows, 4 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/PhysioNet-MI +dataset_kind: window +label_mode: classification + +model: + num_channels: 64 + +model_head: + num_classes: 4 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 4 diff --git a/config/dataset/seed-v.yaml b/config/dataset/seed-v.yaml new file mode 100644 index 0000000..3cd31bb --- /dev/null +++ b/config/dataset/seed-v.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# SEED-V: emotion. 62 channels, 4 s windows, 5 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/SEED-V +dataset_kind: window +label_mode: classification + +model: + num_channels: 62 + +model_head: + num_classes: 5 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 4 diff --git a/config/dataset/seed-vig.yaml b/config/dataset/seed-vig.yaml new file mode 100644 index 0000000..4bc322f --- /dev/null +++ b/config/dataset/seed-vig.yaml @@ -0,0 +1,36 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# SEED-VIG: vigilance regression. 17 channels, 8 s windows, continuous target in [0, 1]. +# MlpRegressionHead takes neither num_classes nor num_patches, so neither is set here. +defaults: + - /model_head: mlp_regression_head + - override /task: finetune_regression_task_SCEReBrO + - override /criterion: mse_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/SEED-VIG +dataset_kind: window +label_mode: regression + +model: + num_channels: 17 + +model_head: + # PERCLOS is bounded, so the head keeps its sigmoid. Disable for unbounded targets. + bounded_output: True diff --git a/config/dataset/shu-mi.yaml b/config/dataset/shu-mi.yaml new file mode 100644 index 0000000..880ba22 --- /dev/null +++ b/config/dataset/shu-mi.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# SHU-MI: motor imagery. 32 channels, 4 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/SHU-MI +dataset_kind: window +label_mode: classification + +model: + num_channels: 32 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 4 diff --git a/config/dataset/stew.yaml b/config/dataset/stew.yaml new file mode 100644 index 0000000..ca68496 --- /dev/null +++ b/config/dataset/stew.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# STEW: workload. 14 channels, 4 s windows, 3 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/STEW +dataset_kind: window +label_mode: classification + +model: + num_channels: 14 + +model_head: + num_classes: 3 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 4 diff --git a/config/dataset/tuab.yaml b/config/dataset/tuab.yaml new file mode 100644 index 0000000..c792865 --- /dev/null +++ b/config/dataset/tuab.yaml @@ -0,0 +1,37 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# TUAB: abnormal vs normal. 22 channels, 10 s windows, 2 classes. +defaults: + - /model_head: mlp_classification_head + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + +dataset_root: ${env:DATA_PATH}/finetuning/TUAB +dataset_kind: window +label_mode: classification + +model: + num_channels: 22 + +model_head: + num_classes: 2 + # One patch per second. Only read when pooling_method is 'flatten'; mean pooling + # takes the window length from the data. + num_patches: 10 diff --git a/config/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml new file mode 100644 index 0000000..42b3a04 --- /dev/null +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -0,0 +1,112 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +tag: SCEReBrO_finetune + +model_size: tiny + +# Download root for the Hugging Face release, so a run can name the checkpoint without +# restating the path: pretrained_safetensors_path='${pretrained_root}/S_CEReBrO.safetensors' +pretrained_root: ${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO + +# Names the registry entry whose BatchRequirements the task enforces on every batch. +model_family: s-cerebro + +# Each file in config/dataset/ owns everything that varies per corpus, including the +# matching head, task and criterion. Switch with: dataset=isruc + +gpus: -1 +num_nodes: 1 +num_workers: 4 +batch_size: 64 + +seed: 0 +training: True +resume: False +final_validate: True +final_test: True +find_unused_parameters: True + +pretrained_checkpoint_path: null +pretrained_safetensors_path: null + +# S-CEReBrO normalises with min-max in LMDBDataset, not here. Disabled so the +# inherited quantile settings cannot be misread as active. +input_normalization: + normalize: False + +defaults: + # model_head, task and criterion come from the dataset group. Setting them here + # would win over it, since Hydra applies a config's own values after its defaults. + - /dataset: tuab + - override /data_module: finetune_data_module_SCEReBrO + - override /model: SCEReBrO_tiny + - override /scheduler: cosine + +task: + freeze_backbone: False # True gives a linear-probe style run + layerwise_lr_decay: 0.95 + +io: + checkpoint_dirpath: ${env:CHECKPOINT_DIR}/checkpoints + base_output_path: ${env:CHECKPOINT_DIR}/outputs + version: ${tag}_${model_size} + +callbacks: + lr_monitor: + _target_: 'pytorch_lightning.callbacks.LearningRateMonitor' + logging_interval: step + progress_bar: + _target_: 'pytorch_lightning.callbacks.TQDMProgressBar' + refresh_rate: 50 + early_stopping: + _target_: 'pytorch_lightning.callbacks.EarlyStopping' + monitor: 'val_loss' + patience: 35 + mode: 'min' + verbose: True + +model_checkpoint: + save_last: True + monitor: "val_loss" + mode: "min" + save_top_k: 1 + every_n_epochs: 1 + +trainer: + accelerator: gpu + num_nodes: ${num_nodes} + devices: ${gpus} + strategy: ddp + max_epochs: 50 + gradient_clip_val: 1 + check_val_every_n_epoch: 1 + accumulate_grad_batches: 1 + +optimizer: + optim: AdamW + lr: 1e-4 + betas: [0.9, 0.999] + weight_decay: 0.05 + +scheduler: + trainer: ${trainer} + warmup_epochs: 5 + min_lr: 1e-5 + warmup_lr_init: 1e-5 diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml new file mode 100644 index 0000000..7a76a88 --- /dev/null +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -0,0 +1,103 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +tag: SCEReBrO_pretrain + +model_size: tiny +model_family: s-cerebro + +gpus: -1 +num_nodes: 1 +num_workers: 8 +batch_size: 32 + +seed: 42 +training: True +resume: False +final_validate: True +final_test: False +find_unused_parameters: False + +pretrained_checkpoint_path: null +pretrained_safetensors_path: null + +# S-CEReBrO normalises with min-max in LMDBDataset, not here. Disabled so the +# inherited quantile settings cannot be misread as active. +input_normalization: + normalize: False + +defaults: + # model_head is absent from the global defaults, so it is added, not overridden. + # Hydra requires additions before overrides. + - /model_head: patch_reconstruction_head + - override /data_module: pretrain_data_module_SCEReBrO + - override /model: SCEReBrO_tiny + - override /task: pretrain_task_SCEReBrO + - override /criterion: masked_reconstruction_loss + - override /scheduler: cosine + +model: + num_channels: 64 + +task: + masking_ratio: 0.5 + +criterion: + loss_type: l2 + alpha: 0.1 + +io: + checkpoint_dirpath: ${env:CHECKPOINT_DIR}/checkpoints + base_output_path: ${env:CHECKPOINT_DIR}/outputs + version: ${tag}_${model_size} + +callbacks: + lr_monitor: + _target_: 'pytorch_lightning.callbacks.LearningRateMonitor' + logging_interval: step + progress_bar: + _target_: 'pytorch_lightning.callbacks.TQDMProgressBar' + refresh_rate: 10 + +model_checkpoint: + save_last: True + monitor: "val_loss" + mode: "min" + save_top_k: 5 + +trainer: + accelerator: gpu + num_nodes: ${num_nodes} + devices: ${gpus} + strategy: ddp + max_epochs: 100 + check_val_every_n_epoch: 5 + accumulate_grad_batches: 8 + +optimizer: + optim: AdamW + lr: 1.25e-3 + betas: [0.9, 0.98] + weight_decay: 0.05 + +scheduler: + trainer: ${trainer} + warmup_epochs: 5 + min_lr: 2.5e-7 + warmup_lr_init: 2.5e-7 diff --git a/config/model/SCEReBrO_tiny.yaml b/config/model/SCEReBrO_tiny.yaml new file mode 100644 index 0000000..1868e98 --- /dev/null +++ b/config/model/SCEReBrO_tiny.yaml @@ -0,0 +1,52 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +# The single released size: 2.4M parameters, matching S_CEReBrO.safetensors. +model: + _target_: models.s_cerebro.SCerebroEncoder + + # Tokenisation. patch_size is fixed at 200 by TemporalConvTokenizer, which is one + # second at the 200 Hz sampling rate every S-CEReBrO dataset is resampled to. + patch_size: 200 + num_channels: 64 # channels in the input; override per fine-tuning dataset + embed_dim: 180 + depth: 6 + num_heads: 5 + mlp_ratio: 4.0 + + # windowed-alternating is the published method. alternating and full are the + # ablation baselines from the paper. + attention_type: windowed-alternating + + drop_path: 0.1 + attn_drop: 0.1 + proj_drop: 0.1 + + # Capacity of the positional tables. An encoder pre-trained at this capacity can be + # fine-tuned on fewer channels and shorter recordings without being rebuilt. + max_channels: 64 + max_timesteps: 6000 + + window_size_spatial: 7 + window_size_temporal: 5 + dilation_cycle_spatial: [1, 2, 4] + dilation_cycle_temporal: [1, 2, 4] + shift_cycle_spatial: [-1, 1, -2, 2] + shift_cycle_temporal: [-1, 1, -2, 2] + use_axial_mode: False diff --git a/config/model_head/mlp_classification_head.yaml b/config/model_head/mlp_classification_head.yaml new file mode 100644 index 0000000..c1e51db --- /dev/null +++ b/config/model_head/mlp_classification_head.yaml @@ -0,0 +1,28 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +model_head: + _target_: models.model_heads.mlp_classification_head.MlpClassificationHead + embed_dim: ${model.embed_dim} + num_classes: 2 # override per fine-tuning dataset + pooling_method: mean # 'mean' or 'flatten' + dropout: 0.1 + pooling: True + num_channels: ${model.num_channels} + num_patches: 10 # window length in seconds; only read by 'flatten' diff --git a/config/model_head/mlp_regression_head.yaml b/config/model_head/mlp_regression_head.yaml new file mode 100644 index 0000000..a6b5e20 --- /dev/null +++ b/config/model_head/mlp_regression_head.yaml @@ -0,0 +1,25 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +model_head: + _target_: models.model_heads.mlp_regression_head.MlpRegressionHead + embed_dim: ${model.embed_dim} + dropout: 0.1 + pooling: True + bounded_output: True # sigmoid; disable for unbounded targets diff --git a/config/model_head/patch_reconstruction_head.yaml b/config/model_head/patch_reconstruction_head.yaml new file mode 100644 index 0000000..0bb1f5b --- /dev/null +++ b/config/model_head/patch_reconstruction_head.yaml @@ -0,0 +1,23 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +model_head: + _target_: models.model_heads.patch_reconstruction_head.PatchReconstructionHead + embed_dim: ${model.embed_dim} + patch_size: ${model.patch_size} diff --git a/config/model_head/sequence_classification_head.yaml b/config/model_head/sequence_classification_head.yaml new file mode 100644 index 0000000..a1b5222 --- /dev/null +++ b/config/model_head/sequence_classification_head.yaml @@ -0,0 +1,32 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +model_head: + _target_: models.model_heads.sequence_classification_head.SequenceClassificationHead + sequence_length: 20 + num_channels: ${model.num_channels} + num_patches: 30 + embed_dim: ${model.embed_dim} + num_classes: 5 + hidden_dim: 512 + num_layers: 1 + nhead: 4 + dim_feedforward: 2048 + dropout: 0.1 + norm_first: True diff --git a/config/task/finetune_regression_task_SCEReBrO.yaml b/config/task/finetune_regression_task_SCEReBrO.yaml new file mode 100644 index 0000000..4fd24a7 --- /dev/null +++ b/config/task/finetune_regression_task_SCEReBrO.yaml @@ -0,0 +1,23 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +task: + _target_: 'tasks.regression_task.RegressionTask' + freeze_backbone: False + layerwise_lr_decay: 0.95 diff --git a/config/task/finetune_task_SCEReBrO.yaml b/config/task/finetune_task_SCEReBrO.yaml new file mode 100644 index 0000000..76cb026 --- /dev/null +++ b/config/task/finetune_task_SCEReBrO.yaml @@ -0,0 +1,24 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +task: + _target_: 'tasks.classification_task.ClassificationTask' + freeze_backbone: False + layerwise_lr_decay: 0.95 + head_lr_multiplier: 1.0 diff --git a/config/task/pretrain_task_SCEReBrO.yaml b/config/task/pretrain_task_SCEReBrO.yaml new file mode 100644 index 0000000..1be7d03 --- /dev/null +++ b/config/task/pretrain_task_SCEReBrO.yaml @@ -0,0 +1,22 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +task: + _target_: 'tasks.mae_pretraining.MaskedAutoencoderPretrainingTask' + masking_ratio: 0.5 diff --git a/criterion/ce_criterion.py b/criterion/ce_criterion.py new file mode 100644 index 0000000..dd88412 --- /dev/null +++ b/criterion/ce_criterion.py @@ -0,0 +1,54 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Dict, Optional, Sequence + +import torch +from torch import nn + + +class CrossEntropyLossWrapper(nn.Module): + """Cross-entropy loss with optional label smoothing and class weighting. + + Class weights are held by the wrapped ``nn.CrossEntropyLoss`` as a buffer, so + they follow the enclosing task onto whichever device it is moved to. + + Args: + label_smoothing: Smoothing coefficient in ``[0, 1)``. + weight: Optional per-class weights, ordered by class index. + """ + + def __init__(self, label_smoothing: float = 0.0, weight: Optional[Sequence[float]] = None): + super().__init__() + self.label_smoothing = label_smoothing + class_weight = torch.tensor(weight, dtype=torch.float32) if weight is not None else None + self.loss_fn = nn.CrossEntropyLoss(label_smoothing=label_smoothing, weight=class_weight) + + def forward(self, pred: torch.Tensor, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + """Compute cross-entropy against ``batch['label']``. + + Args: + pred: Logits of shape ``(batch, num_classes)``. + batch: Must contain ``label`` with integer class indices of shape ``(batch,)``. + + Returns: + Scalar loss. + """ + return self.loss_fn(pred, batch["label"]) diff --git a/criterion/focal_criterion.py b/criterion/focal_criterion.py new file mode 100644 index 0000000..0bf154b --- /dev/null +++ b/criterion/focal_criterion.py @@ -0,0 +1,65 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Dict + +import torch +import torch.nn.functional as F +from torch import nn + + +class FocalLossWrapper(nn.Module): + """Binary focal loss for strongly imbalanced classification. + + Down-weights well-classified examples so that rare positives dominate the + gradient, which is the regime of seizure detection (CHB-MIT, Neonate). The loss + is evaluated on the positive-class logit and is computed through + ``logsigmoid`` so that confident predictions cannot produce infinities. + + Args: + alpha: Weight of the positive class in ``[0, 1]``. + gamma: Focusing exponent; larger values suppress easy examples more. + """ + + def __init__(self, alpha: float = 0.8, gamma: float = 0.7): + super().__init__() + self.alpha = float(alpha) + self.gamma = float(gamma) + + def forward(self, pred: torch.Tensor, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + """Compute focal loss against ``batch['label']``. + + Args: + pred: Logits of shape ``(batch, 2)`` or ``(batch,)``. + batch: Must contain ``label`` with binary targets of shape ``(batch,)``. + + Returns: + Scalar loss. + """ + # The class-1 logit, matching the BIOT reference this follows. The loss is + # therefore not invariant to a constant shift of both logits; the published + # results were produced this way, so do not substitute the class margin. + logits = pred[:, 1] if pred.dim() == 2 and pred.shape[1] == 2 else pred.reshape(-1) + targets = batch["label"].reshape(-1).to(logits.dtype) + + prob = torch.sigmoid(logits) + positive = -self.alpha * (1 - prob) ** self.gamma * targets * F.logsigmoid(logits) + negative = -(1 - self.alpha) * prob ** self.gamma * (1 - targets) * F.logsigmoid(-logits) + return (positive + negative).mean() diff --git a/criterion/masked_reconstruction_loss.py b/criterion/masked_reconstruction_loss.py new file mode 100644 index 0000000..c666e5b --- /dev/null +++ b/criterion/masked_reconstruction_loss.py @@ -0,0 +1,92 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Dict, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + + +class MaskedReconstructionLoss(nn.Module): + """Reconstruction loss for masked EEG pre-training. + + The loss is averaged over masked, non-padded patches. When ``alpha`` is + non-zero, the loss on visible non-padded patches is added with weight ``alpha``, + which stabilises early training without letting the visible patches dominate. + + Args: + loss_type: One of ``l1``, ``l2``, ``smooth_l1``. + alpha: Weight of the visible-patch term; ``0`` disables it. + """ + + def __init__(self, loss_type: str = "l2", alpha: float = 0.1): + super().__init__() + if loss_type not in {"l1", "l2", "smooth_l1"}: + raise ValueError(f"loss_type must be 'l1', 'l2' or 'smooth_l1', got '{loss_type}'") + self.loss_type = loss_type + self.alpha = alpha + + def _elementwise_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + if self.loss_type == "l1": + return F.l1_loss(pred, target, reduction="none") + if self.loss_type == "l2": + return F.mse_loss(pred, target, reduction="none") + return F.smooth_l1_loss(pred, target, reduction="none") + + def forward(self, pred: torch.Tensor, batch: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, Dict[str, float]]: + """Compute the masked reconstruction loss. + + Args: + pred: Reconstructed patches of shape ``(batch, num_tokens, patch_size)``. + batch: Must contain ``target`` with the same shape as ``pred``, ``token_mask`` + of shape ``(batch, num_tokens)`` where 1 marks a masked token, and + ``attn_mask`` of shape ``(batch, num_tokens)`` where 1 marks a real token. + + Returns: + Tuple of the scalar loss and a dictionary of values to log. + """ + target = batch["target"] + token_mask = batch["token_mask"].bool() + attn_mask = batch.get("attn_mask") + attn_mask = torch.ones_like(token_mask) if attn_mask is None else attn_mask.bool() + + per_patch = self._elementwise_loss(pred, target).mean(dim=-1) + + masked = self._mean_or_zero(per_patch, token_mask & attn_mask) + logs = {"masked_loss": masked.detach()} + + if self.alpha == 0: + return masked, logs + + visible = self._mean_or_zero(per_patch, (~token_mask) & attn_mask) + logs["visible_loss"] = visible.detach() + return masked + self.alpha * visible, logs + + @staticmethod + def _mean_or_zero(values: torch.Tensor, selection: torch.Tensor) -> torch.Tensor: + """Mean of the selected entries, or a differentiable zero if none are selected. + + ``values[selection].mean()`` is NaN on an empty selection, which would poison + every gradient. Branching on ``selection.any()`` would avoid that but synchronise + with the GPU on every step, so the empty case is handled arithmetically instead. + """ + selected = torch.where(selection, values, values.new_zeros(())) + return selected.sum() / selection.sum().clamp(min=1).to(values.dtype) diff --git a/criterion/mse_criterion.py b/criterion/mse_criterion.py new file mode 100644 index 0000000..d6f7e3a --- /dev/null +++ b/criterion/mse_criterion.py @@ -0,0 +1,45 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Dict + +import torch +from torch import nn + + +class MSELossWrapper(nn.Module): + """Mean squared error loss for scalar regression targets.""" + + def __init__(self, reduction: str = "mean"): + super().__init__() + self.loss_fn = nn.MSELoss(reduction=reduction) + + def forward(self, pred: torch.Tensor, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + """Compute MSE against ``batch['label']``. + + Args: + pred: Predictions broadcastable to the target shape. + batch: Must contain ``label`` with the regression targets. + + Returns: + Scalar loss. + """ + target = batch["label"] + return self.loss_fn(pred.reshape(target.shape).to(target.dtype), target) diff --git a/data_module/finetuning_data_module.py b/data_module/finetuning_data_module.py new file mode 100644 index 0000000..04c100c --- /dev/null +++ b/data_module/finetuning_data_module.py @@ -0,0 +1,83 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Optional + +import pytorch_lightning as pl +from torch.utils.data import DataLoader, Dataset + + +class FinetuningDataModule(pl.LightningDataModule): + """Data module for fine-tuning on a dataset with fixed train, validation and test splits. + + The three splits are supplied as already-constructed datasets, so the split is + defined entirely by the preprocessing step and is identical for every run and every + model. No resampling, subject-level cross-validation or automatic splitting happens + here. + + Args: + train: Training dataset. + val: Validation dataset, used for model selection and early stopping. + test: Test dataset, evaluated once at the end of the run. + cfg: Holds ``batch_size`` and ``num_workers``. + name: Label for logging. + """ + + def __init__( + self, + train: Dataset, + val: Dataset, + test: Optional[Dataset] = None, + cfg=None, + name: str = "", + **kwargs, + ): + super().__init__() + self.train_dataset = train + self.val_dataset = val + self.test_dataset = test + self.cfg = cfg + self.name = name + + def _loader(self, dataset: Dataset, shuffle: bool, drop_last: bool) -> DataLoader: + """Build a loader with the shared worker and pinning settings.""" + if dataset is None: + raise ValueError("Requested a dataloader for a split that was not configured") + return DataLoader( + dataset, + batch_size=self.cfg.batch_size, + shuffle=shuffle, + num_workers=self.cfg.num_workers, + drop_last=drop_last, + pin_memory=True, + persistent_workers=self.cfg.num_workers > 0, + ) + + def train_dataloader(self) -> DataLoader: + """Shuffled loader over the training split, dropping the last partial batch.""" + return self._loader(self.train_dataset, shuffle=True, drop_last=True) + + def val_dataloader(self) -> DataLoader: + """Deterministic loader over the full validation split.""" + return self._loader(self.val_dataset, shuffle=False, drop_last=False) + + def test_dataloader(self) -> DataLoader: + """Deterministic loader over the full test split.""" + return self._loader(self.test_dataset, shuffle=False, drop_last=False) diff --git a/data_module/pretraining_data_module.py b/data_module/pretraining_data_module.py new file mode 100644 index 0000000..95231c3 --- /dev/null +++ b/data_module/pretraining_data_module.py @@ -0,0 +1,97 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +from typing import Dict, Optional + +import pytorch_lightning as pl +import torch +from torch.utils.data import ConcatDataset, DataLoader, Dataset + + +class PretrainingDataModule(pl.LightningDataModule): + """Data module for self-supervised pre-training on a union of EEG corpora. + + The configured corpora are concatenated and split once into a training and a + validation portion with a seeded generator, so the split is identical across + processes in distributed training and reproducible across runs. The split is over + windows rather than recordings, so the validation loss measures reconstruction + quality in distribution rather than generalisation to unseen subjects. + + All corpora must share a window length in timesteps, otherwise batches drawn + across corpora cannot be collated. + + Args: + train: Mapping of corpus name to dataset; ``None`` entries are skipped. + cfg: Holds ``batch_size`` and ``num_workers``. + val_ratio: Fraction of windows held out for validation. + seed: Seed for the train/validation split. + name: Label for logging. + """ + + def __init__( + self, + train: Dict[str, Optional[Dataset]], + cfg=None, + val_ratio: float = 0.2, + seed: int = 42, + name: str = "", + **kwargs, + ): + super().__init__() + datasets = [dataset for dataset in train.values() if dataset is not None] + if not datasets: + raise ValueError("No pre-training datasets were configured") + + combined = ConcatDataset(datasets) + generator = torch.Generator().manual_seed(seed) + self.train_dataset, self.val_dataset = torch.utils.data.random_split( + combined, [1.0 - val_ratio, val_ratio], generator=generator + ) + print( + f"[PretrainingDataModule] corpora={len(datasets)} windows={len(combined)} " + f"train={len(self.train_dataset)} val={len(self.val_dataset)}" + ) + + self.cfg = cfg + self.name = name + + def train_dataloader(self) -> DataLoader: + """Shuffled loader over the training split.""" + return DataLoader( + self.train_dataset, + batch_size=self.cfg.batch_size, + shuffle=True, + num_workers=self.cfg.num_workers, + drop_last=True, + pin_memory=True, + persistent_workers=self.cfg.num_workers > 0, + ) + + def val_dataloader(self) -> DataLoader: + """Deterministic loader over the validation split.""" + return DataLoader( + self.val_dataset, + batch_size=self.cfg.batch_size, + shuffle=False, + num_workers=self.cfg.num_workers, + drop_last=False, + pin_memory=True, + persistent_workers=self.cfg.num_workers > 0, + ) diff --git a/datasets/lmdb_dataset.py b/datasets/lmdb_dataset.py new file mode 100644 index 0000000..d201235 --- /dev/null +++ b/datasets/lmdb_dataset.py @@ -0,0 +1,298 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +import glob +import json +import os +import pickle +from collections import OrderedDict +from typing import Any, Dict, List, Optional, Tuple + +import lmdb +import numpy as np +import torch +import torch.nn.functional as F + +# LMDB refuses to open a path twice in one process, and run_train.py re-instantiates the +# data module before testing while the first one may still hold its environment. Every +# reader in the process therefore shares one reference-counted environment per path. +_SHARED_ENVS: Dict[str, List[Any]] = {} + + +def acquire_env(path: str) -> Any: + """Return the process's read-only environment for ``path``, opening it if needed.""" + key = os.path.realpath(path) + entry = _SHARED_ENVS.get(key) + if entry is None: + env = lmdb.open(path, readonly=True, lock=False, readahead=True, max_readers=64, map_async=True) + entry = _SHARED_ENVS[key] = [env, 0] + entry[1] += 1 + return entry[0] + + +def release_env(path: str) -> None: + """Drop one reference to ``path``'s environment, closing it after the last one.""" + key = os.path.realpath(path) + entry = _SHARED_ENVS.get(key) + if entry is None: + return + entry[1] -= 1 + if entry[1] <= 0: + del _SHARED_ENVS[key] + entry[0].close() + + +def resolve_lmdb_path(path: str, split: Optional[str] = None) -> Tuple[str, Optional[str]]: + """Resolve a dataset directory and split name to an LMDB file and its index file. + + Args: + path: Either a directory containing ``{split}.lmdb`` files or a path ending in + ``.lmdb``. + split: One of ``train``, ``val``, ``test``, or ``None`` to auto-select. + + Returns: + Tuple of the LMDB path and the matching ``{stem}.meta.json`` path, or ``None`` + when no index file exists. A split-specific index is never substituted with a + pooled one, so a missing ``train.meta.json`` falls back to scanning + ``train.lmdb`` rather than silently indexing it with pooled keys. + """ + if path.endswith(".lmdb"): + stem = os.path.splitext(os.path.basename(path))[0] + meta = os.path.join(os.path.dirname(path), f"{stem}.meta.json") + return path, meta if os.path.isfile(meta) else None + + if not os.path.isdir(path): + raise ValueError(f"Dataset path is neither a directory nor an .lmdb file: {path}") + + available = { + os.path.splitext(os.path.basename(p))[0]: p for p in glob.glob(os.path.join(path, "*.lmdb")) + } + if not available: + raise FileNotFoundError(f"No .lmdb files found in {path}") + + if split in {"train", "val", "test"}: + if split not in available: + raise FileNotFoundError( + f"Expected {split}.lmdb in {path}; found {sorted(available)}. " + "Re-run the preprocessing script for this dataset to build the splits." + ) + stem = split + elif len(available) == 1: + stem = next(iter(available)) + else: + raise ValueError(f"Multiple LMDBs in {path}; set split to one of {sorted(available)}") + + meta = os.path.join(path, f"{stem}.meta.json") + return available[stem], meta if os.path.isfile(meta) else None + + +class LMDBDataset(torch.utils.data.Dataset): + """LMDB-backed EEG dataset for windowed and sequence samples. + + Each LMDB value is a pickled dictionary with keys ``eeg``, ``channel_coords`` and + optionally ``label`` and ``subject_id``. + + Two sample layouts are supported: + + * ``window``: ``eeg`` has shape ``(num_channels, num_timesteps)`` with one scalar label. + * ``sequence``: ``eeg`` has shape ``(sequence_length, num_channels, num_timesteps)`` + with one label per element. Every entry in a sequence dataset must share the same + ``sequence_length``, otherwise batches cannot be collated. + + Args: + path: Dataset directory or ``.lmdb`` file. + dataset_kind: ``window`` or ``sequence``. + split: ``train``, ``val`` or ``test``. + apply_minmax: Scale each channel to ``[-1, 1]`` independently. + apply_zero_padding: Zero-pad channels and timesteps up to ``max_channels`` and + ``max_timesteps``, reporting how much padding was added. + max_channels: Channel count to pad to; required if ``apply_zero_padding``. + max_timesteps: Timestep count to pad to; optional. + label_mode: ``classification``, ``regression``, or ``auto`` to infer from dtype. + use_cache: Keep decoded samples in an in-process LRU cache. + cache_size: Maximum number of cached samples. + """ + + def __init__( + self, + path: str, + dataset_kind: str = "window", + split: Optional[str] = None, + apply_minmax: bool = True, + apply_zero_padding: bool = False, + max_channels: Optional[int] = None, + max_timesteps: Optional[int] = None, + label_mode: str = "auto", + use_cache: bool = False, + cache_size: int = 1000, + ): + if dataset_kind not in {"window", "sequence"}: + raise ValueError(f"dataset_kind must be 'window' or 'sequence', got '{dataset_kind}'") + if label_mode not in {"auto", "classification", "regression"}: + raise ValueError(f"Invalid label_mode '{label_mode}'") + + self.dataset_kind = dataset_kind + self.apply_minmax = apply_minmax + self.apply_zero_padding = apply_zero_padding + self.max_channels = max_channels + self.max_timesteps = max_timesteps + self.label_mode = label_mode + self.use_cache = use_cache + self.cache_size = cache_size + + self.lmdb_path, self.meta_path = resolve_lmdb_path(path, split) + self.env = None + self.keys = self._load_keys() + self.cache: "OrderedDict[int, Dict[str, Any]]" = OrderedDict() + + print(f"[LMDBDataset] {os.path.basename(self.lmdb_path)}: {len(self.keys)} entries") + + def _ensure_env(self) -> None: + """Attach to the LMDB environment on first read. + + Opening lazily keeps an open handle out of forked dataloader workers, which + cannot open a path their parent already had open. + """ + if self.env is None: + self.env = acquire_env(self.lmdb_path) + + def _load_keys(self) -> List[bytes]: + """Read the sample keys from the index file, or scan the LMDB if there is none.""" + if self.meta_path is not None: + with open(self.meta_path, "r") as handle: + meta = json.load(handle) + keys = meta.get("keys", []) + if not keys: + raise ValueError(f"Index file {self.meta_path} contains no keys") + return [key.encode() if isinstance(key, str) else key for key in keys] + + env = acquire_env(self.lmdb_path) + try: + with env.begin(buffers=True) as txn: + return [bytes(key) for key in txn.cursor().iternext(keys=True, values=False)] + finally: + release_env(self.lmdb_path) + + def __len__(self) -> int: + """Number of samples.""" + return len(self.keys) + + def _min_max_normalize(self, x: torch.Tensor) -> torch.Tensor: + """Scale each channel independently to ``[-1, 1]``.""" + flat = x.reshape(-1, x.shape[-1]) + minimum = flat.min(dim=1, keepdim=True).values + maximum = flat.max(dim=1, keepdim=True).values + # Absolute epsilon, as in the reference that produced the released checkpoint. + # Negligible for corpora stored in µV; for those stored in volts it is around 1% + # of a typical range, so their windows land slightly inside [-1, 1]. + scaled = (flat - minimum) / (maximum - minimum + 1e-6) + return ((scaled - 0.5) * 2).reshape(x.shape) + + def _pad(self, x: torch.Tensor, channel_coords: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, int, int]: + """Zero-pad the channel and time axes, returning the amount of padding added.""" + channels, timesteps = x.shape[-2], x.shape[-1] + pad_time = max(0, (self.max_timesteps or timesteps) - timesteps) + pad_channels = max(0, (self.max_channels or channels) - channels) + padded = F.pad(x, (0, pad_time, 0, pad_channels), value=0.0) + + coord_padding = max(0, (self.max_channels or channel_coords.shape[0]) - channel_coords.shape[0]) + padded_coords = F.pad(channel_coords, (0, 0, 0, 0, 0, coord_padding), value=0.0) + return padded, padded_coords, pad_channels, pad_time + + def _is_regression(self, label: Any) -> bool: + """Decide whether a stored label is a regression target.""" + if self.label_mode != "auto": + return self.label_mode == "regression" + if isinstance(label, (float, np.floating)): + return True + if isinstance(label, (list, np.ndarray)): + array = np.asarray(label) + return array.size > 0 and np.issubdtype(array.dtype, np.floating) + return False + + def _build_label(self, label: Any, is_regression: bool) -> torch.Tensor: + """Convert a stored label into a tensor with the right dtype and shape.""" + dtype = torch.float32 if is_regression else torch.long + if self.dataset_kind == "sequence": + array = np.asarray(label, dtype=np.float64 if is_regression else np.int64) + return torch.as_tensor(array.reshape(-1), dtype=dtype) + return torch.tensor(float(label) if is_regression else int(label), dtype=dtype) + + def __getitem__(self, idx: int) -> Dict[str, Any]: + """Decode, normalise and pad one sample.""" + if self.use_cache and idx in self.cache: + self.cache.move_to_end(idx) + return self.cache[idx] + + self._ensure_env() + key = self.keys[idx] + with self.env.begin() as txn: + raw = txn.get(key) + if raw is None: + raise KeyError( + f"Key {key!r} is missing from {self.lmdb_path}. The index file and the " + "LMDB are out of sync; rebuild the dataset." + ) + entry = pickle.loads(raw) + + x = torch.as_tensor(entry["eeg"], dtype=torch.float32) + channel_coords = torch.as_tensor(entry["channel_coords"], dtype=torch.float32) + + expected_dims = 2 if self.dataset_kind == "window" else 3 + if x.dim() != expected_dims: + raise ValueError( + f"dataset_kind='{self.dataset_kind}' expects {expected_dims}-dimensional eeg, " + f"got shape {tuple(x.shape)}" + ) + + if self.apply_minmax: + x = self._min_max_normalize(x) + + pad_channels, pad_time = 0, 0 + if self.apply_zero_padding: + x, channel_coords, pad_channels, pad_time = self._pad(x, channel_coords) + + sample: Dict[str, Any] = { + "input": x, + "channel_coords": channel_coords, + "num_padded_channels": pad_channels, + "num_padded_timesteps": pad_time, + } + + raw_label = entry.get("label") + if raw_label is not None: + is_regression = self._is_regression(raw_label) + sample["label"] = self._build_label(raw_label, is_regression) + + if self.use_cache: + self.cache[idx] = sample + if len(self.cache) > self.cache_size: + self.cache.popitem(last=False) + + return sample + + def __del__(self): + """Release this dataset's reference to the LMDB environment.""" + if getattr(self, "env", None) is not None: + self.env = None + try: + release_env(self.lmdb_path) + except Exception: + pass diff --git a/datasets/tueg_dataset.py b/datasets/tueg_dataset.py new file mode 100644 index 0000000..255ac49 --- /dev/null +++ b/datasets/tueg_dataset.py @@ -0,0 +1,144 @@ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: BioFoundation Contributors * +#* * +#*----------------------------------------------------------------------------* + +import os +from collections import OrderedDict +from typing import Any, Dict, List, Optional + +import numpy as np +import torch +from torch.utils.data import Dataset + +from datasets.lmdb_dataset import acquire_env, release_env + + +class TUEGDataset(Dataset): + """LMDB-backed reader for the TUEG pre-training corpus. + + TUEG is large enough that per-sample pickling is a measurable overhead, so each + value is stored as a fixed-size raw byte blob holding the padded waveform followed + by the electrode coordinates. Sample keys live in a companion text file, one per + line, which avoids a full LMDB scan at start-up. + + Recordings are padded to ``max_channels`` offline. The number of padded channels is + recovered by finding channels whose coordinates are all zero, so pre-training can + exclude them from masking and from attention. + + Args: + lmdb_path: Path to the LMDB directory. + keys_path: Path to the newline-separated key list; defaults to + ``{lmdb_path without extension}_keys.txt``. + max_channels: Channel count every record was padded to. + sampling_freq: Sampling frequency in Hz. + slice_duration: Window length in seconds. + use_cache: Keep decoded samples in an in-process LRU cache. + cache_size: Maximum number of cached samples. + """ + + def __init__( + self, + lmdb_path: str, + keys_path: Optional[str] = None, + max_channels: int = 64, + sampling_freq: int = 200, + slice_duration: int = 30, + use_cache: bool = False, + cache_size: int = 10000, + ): + self.lmdb_path = lmdb_path + self.keys_path = keys_path or f"{os.path.splitext(lmdb_path)[0]}_keys.txt" + self.max_channels = max_channels + self.num_timesteps = sampling_freq * slice_duration + + self.waveform_bytes = max_channels * self.num_timesteps * np.dtype(np.float32).itemsize + self.coords_bytes = max_channels * 2 * 3 * np.dtype(np.float32).itemsize + self.record_bytes = self.waveform_bytes + self.coords_bytes + + if not os.path.isfile(self.keys_path): + raise FileNotFoundError(f"Key list not found: {self.keys_path}") + with open(self.keys_path, "r") as handle: + self.keys: List[bytes] = [line.strip().encode("ascii") for line in handle if line.strip()] + + self.use_cache = use_cache + self.cache_size = cache_size + self.cache: "OrderedDict[int, Dict[str, Any]]" = OrderedDict() + self.env = None + + print(f"[TUEGDataset] {os.path.basename(self.lmdb_path)}: {len(self.keys)} windows") + + def _ensure_env(self) -> None: + """Attach to the LMDB environment lazily so it is not inherited across worker forks.""" + if self.env is None: + self.env = acquire_env(self.lmdb_path) + + def __len__(self) -> int: + """Number of windows.""" + return len(self.keys) + + def __getitem__(self, idx: int) -> Dict[str, Any]: + """Decode one window and its electrode coordinates.""" + if self.use_cache and idx in self.cache: + self.cache.move_to_end(idx) + return self.cache[idx] + + self._ensure_env() + key = self.keys[idx] + with self.env.begin() as txn: + blob = txn.get(key) + + if blob is None: + raise KeyError(f"Key {key!r} is missing from {self.lmdb_path}") + if len(blob) != self.record_bytes: + raise ValueError( + f"Record {key!r} has {len(blob)} bytes, expected {self.record_bytes}. " + "Check that max_channels, sampling_freq and slice_duration match the " + "values used during preprocessing." + ) + + waveform = np.frombuffer(blob[: self.waveform_bytes], dtype=np.float32) + waveform = waveform.reshape(self.max_channels, self.num_timesteps).copy() + coords = np.frombuffer(blob[self.waveform_bytes :], dtype=np.float32) + coords = coords.reshape(self.max_channels, 2, 3).copy() + + channel_coords = torch.from_numpy(coords) + num_padded = int((channel_coords.view(self.max_channels, -1).abs().sum(dim=1) == 0).sum()) + + sample = { + "input": torch.from_numpy(waveform), + "channel_coords": channel_coords, + "num_padded_channels": num_padded, + "num_padded_timesteps": 0, + } + + if self.use_cache: + self.cache[idx] = sample + if len(self.cache) > self.cache_size: + self.cache.popitem(last=False) + + return sample + + def __del__(self): + """Release this dataset's reference to the LMDB environment.""" + if getattr(self, "env", None) is not None: + self.env = None + try: + release_env(self.lmdb_path) + except Exception: + pass diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md index 630f1ef..4c368fe 100644 --- a/docs/CITATIONS.md +++ b/docs/CITATIONS.md @@ -11,6 +11,7 @@ Please cite the paper corresponding to the model used in your work. Published pr | TinyMyo | [arXiv preprint](https://arxiv.org/abs/2512.15729). | | LuMamba | Accepted at EUSIPCO 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2603.19100)). | | PanLUNA | Accepted at IEEE AICAS 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2604.04297)). | +| S-CEReBrO | Accepted at MICCAI 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2607.27913)). | ## FEMBA @@ -85,4 +86,20 @@ Please cite the paper corresponding to the model used in your work. Published pr } ``` -The LuMamba and PanLUNA entries should be updated with their final DOI and page numbers once the EUSIPCO and AICAS proceedings are published. +## S-CEReBrO + +```bibtex +@inproceedings{bucagu2026scerebro, + title={{S-CEReBrO}: Breaking the Memory Barrier in Continuous {EEG} Monitoring}, + author={Bucagu, Glenn Anta and Ingolfsson, Thorir Mar and Li, Yawei and Benini, Luca}, + booktitle={Medical Image Computing and Computer Assisted Intervention (MICCAI)}, + year={2026}, + note={Accepted; proceedings forthcoming}, + eprint={2607.27913}, + archivePrefix={arXiv}, + primaryClass={eess.SP}, + url={https://arxiv.org/abs/2607.27913} +} +``` + +The LuMamba, PanLUNA and S-CEReBrO entries should be updated with their final DOI and page numbers once the EUSIPCO, AICAS and MICCAI proceedings are published. diff --git a/docs/README.md b/docs/README.md index e2f88ea..87ba8cf 100644 --- a/docs/README.md +++ b/docs/README.md @@ -21,6 +21,7 @@ Architecture and usage notes are available for every published model family: - [TinyMyo](./model/TinyMyo.md): compact foundation model for sEMG. - [LuMamba](./model/LuMamba.md): query-unified Mamba for EEG. - [PanLUNA](./model/PanLUNA.md): sensor-aware modeling across EEG, ECG, and PPG. +- [S-CEReBrO](./model/SCEReBrO.md): windowed alternating attention for EEG, with a separate prediction head. Each model also has a pretrained release linked from the root model zoo. The canonical machine-readable index is [`biofoundation/model_registry.py`](../biofoundation/model_registry.py). @@ -39,7 +40,8 @@ The training guide covers environment variables, Hydra experiment selection, the ### 5. Project References - [`CONTRIBUTING.md`](../CONTRIBUTING.md) defines extension and pull request expectations. -- [`CITATIONS.md`](./CITATIONS.md) contains BibTeX for all five model families. +- [`docs/adr`](./adr/) records architecture decisions affecting the shared contracts. +- [`CITATIONS.md`](./CITATIONS.md) contains BibTeX for all six model families. - [`config/README.md`](../config/README.md) explains Hydra composition and overrides. - [`make_datasets/README.md`](../make_datasets/README.md) documents preprocessing and HDF5 conversion. diff --git a/docs/adr/0001-two-electrode-geometry-representations.md b/docs/adr/0001-two-electrode-geometry-representations.md new file mode 100644 index 0000000..dc973df --- /dev/null +++ b/docs/adr/0001-two-electrode-geometry-representations.md @@ -0,0 +1,31 @@ +Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE at the repository root for details. + +# ADR 0001: Two electrode-geometry representations + +Status: accepted + +## Context + +Model families in this repository describe electrode geometry in two incompatible ways. + +LUNA, LuMamba and PanLUNA consume `channel_locations` of shape `(batch, channels, 3)`: one 3D coordinate per channel. For a bipolar derivation such as `FP1-F7`, [`models/modules/channel_embeddings.py`](../../models/modules/channel_embeddings.py) resolves both electrodes and returns their midpoint. Channels that are not bipolar contribute their scalp position alone, with no representation of the reference. + +S-CEReBrO consumes `channel_coords` of shape `(batch, channels, 2, 3)`: both electrodes of every channel, kept separate. Its channel embedding maps each electrode independently through a shared MLP and concatenates the halves, and it rejects input that does not carry exactly two electrodes per channel. The preprocessing in [`make_datasets/electrode_positions.py`](../../make_datasets/electrode_positions.py) assigns explicit coordinates to references that have no scalp position, so an average-reference or linked-ears channel is representable. + +The two are not interchangeable. `channel_coords` is strictly richer: the midpoint is recoverable from it, but the pair is not recoverable from the midpoint, and the reference is not represented in `channel_locations` at all. + +## Decision + +Both representations are first-class, independent fields on `SignalBatch`. Neither is derived from the other at runtime, and no conversion is applied automatically. + +A model declares which one it consumes through `BatchRequirements` in [`biofoundation/model_registry.py`](../../biofoundation/model_registry.py). A dataset produces whichever its target family requires. `require_batch_fields` validates the declared field and does not accept the other in its place. + +## Consequences + +The five existing families are unaffected. Their datasets, their geometry handling, and their published checkpoints are untouched, and the new field is unreachable from their code paths. + +A dataset prepared for S-CEReBrO cannot be fed to LUNA, LuMamba or PanLUNA without an explicit conversion step, and vice versa. This is accepted. Making `channel_coords` canonical with `channel_locations` derived from it would have enabled that reuse, but it would have required changing how the existing families obtain geometry, which is not worth the risk to reproducibility of published results. + +If cross-family dataset reuse becomes valuable later, the additive way to get it is a documented reduction helper that a dataset or task calls explicitly. That remains available and is not foreclosed by this decision. What is foreclosed is *implicit* conversion inside `require_batch_fields`, because a model that silently receives midpoints where it expected electrode pairs would train without error and produce quietly wrong geometry. + +Adding a third representation is not acceptable without superseding this record. Two exist because two model families genuinely need different information; a third would mean the contract had stopped being designed. diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md new file mode 100644 index 0000000..ca82b5c --- /dev/null +++ b/docs/model/SCEReBrO.md @@ -0,0 +1,274 @@ +Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache-2.0. See LICENSE at the repository root for details. + +## S-CEReBrO + +S-CEReBrO is a compact EEG encoder built on windowed alternating attention. It tokenises a recording at per-channel patch granularity and alternates attention between the channel axis and the time axis, restricting each pass to a dilated, shifted window. Attention cost per block is linear in the token count rather than quadratic, which is what allows a 64-channel, 30-second window to be modelled directly. + +It is the first family in this repository to separate the encoder from its output layer. The encoder emits token embeddings and nothing else; a [prediction head](../../config/model_head/) turns those into a reconstruction, a class label, or a scalar. One pre-trained encoder therefore serves every downstream task without being rebuilt. + +### Default Input Assumptions + +| Property | Value | +| --- | --- | +| Signal | Scalp EEG | +| Sampling rate | 200 Hz | +| Patch size | 200 samples (1 s), fixed by the tokeniser | +| Channels | Up to `max_channels` (64 by default) | +| Window | Up to `max_timesteps` (6000 samples, 30 s) | +| Amplitude | Per-channel min-max scaled to `[-1, 1]` | + +The encoder requires `channel_coords` of shape `(batch, channels, 2, 3)`: the 3D coordinates of **both** electrodes forming each channel. This is the second electrode-geometry representation in the repository and is deliberately distinct from the `channel_locations` used by LUNA, LuMamba and PanLUNA, which carry one midpoint per channel. See [ADR 0001](../adr/0001-two-electrode-geometry-representations.md) for why both exist and neither is derived from the other. + +Keeping both electrodes separate is what lets a bipolar derivation and a scalp-electrode-plus-reference channel stay distinguishable, and it is why the channel embedding is a function of geometry rather than of a channel index. Montages with different channel counts and orderings share the same parameters. + +### Channel Counts + +The published encoders are pre-trained at 64 channels, and one checkpoint fine-tunes onto any montage from 1 to `max_channels` without modification. Nothing in the state dict depends on the channel count: + +- the temporal position table is sized by `max_timesteps // patch_size` and sliced to the patches present, and is shared across channels; +- the channel embedding is an MLP over 3D electrode coordinates, so it has no per-channel parameters and generalises to montages it never saw; +- the attention blocks take the channel count only as a reshape argument. + +Set `model.num_channels` to the montage you are fine-tuning on and leave `max_channels` at the value the checkpoint was pre-trained with: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune model.num_channels=6 +``` + +The encoder validates this rather than guessing: passing input whose channel count differs from `model.num_channels` raises rather than silently reshaping. Set it to match the dataset. + +Two channel counts coexist and mean different things. `model.num_channels` is how many channels the encoder is built for and must equal what the dataset yields. `model.max_channels` is the capacity the positional table was sized at, and must stay at the pre-training value or the checkpoint will not match. + +For pre-training, corpora with fewer channels are zero-padded up to `max_channels` by `LMDBDataset`, and the padded channels are replaced by a learned pad token, excluded from masking, and masked out of attention. Fine-tuning does not pad: the encoder is simply built at the montage's own size. + +This is why the registry entry declares only `channel_coords` and not `num_padded_channels`. Padding metadata is required by pre-training and irrelevant to fine-tuning, and both stages share one registry entry, so declaring it would force every future S-CEReBrO dataset to emit a field half the pipeline ignores. The pre-training task reads it with `batch.get`, which tolerates its absence. + +### Preprocessing + +Datasets are prepared into LMDB with [`make_datasets`](../../make_datasets/). Each entry is a pickled dictionary with `eeg`, `channel_coords`, and optionally `label` and `subject_id`. Electrode coordinates come from [`make_datasets/electrode_positions.py`](../../make_datasets/electrode_positions.py), which takes its angles from the BESA electrode and surface location tables and assigns fixed coordinates to reference electrodes that have no scalp position of their own, such as the average reference and linked ears. BESA defines theta (its "azimuth") as the angle from Cz, positive on the right, and phi (its "latitude") as the angle from T7 or T8, and positions follow BESA's conversion to the unit sphere: x = sin θ cos φ, y = sin θ sin φ, z = cos θ, with x towards the right ear, y towards the nose and z towards the vertex. + +Neonate windows are labelled from the seconds they contain: a window is a seizure if any of its seconds is, and is skipped if any second lacks annotator consensus. The reference implementation gave window *k* the label of second *k* although the window spans seconds 5*k* to 5*k* + 5, so neonate results are not directly comparable with those produced by its preprocessing. + +Every script reads the raw corpus from `--input_dir` and writes to `--output_dir`; neither has a default. The configs read the results from `$DATA_PATH` laid out as below, and `dataset_root=` for a fine-tuning corpus or `pretraining_root=` for all pre-training corpora points them anywhere else. + +| Corpus | Steps, as `python -m make_datasets.