diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 74592ef..a5b9933 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`. Set `model_family` in the experiment so tasks can enforce the declared batch requirements 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. @@ -41,6 +43,14 @@ python -m unittest discover -s tests -p 'test_*.py' -v python -m compileall -q biofoundation run_train.py models tasks datasets data_module ``` +Tests that import PyTorch live under `tests/model_tests/` and run with `pytest`. That +directory is not a package, so `unittest discover` does not traverse it and the fast +suite stays runnable without the training dependencies installed. + +```bash +pytest tests/model_tests -v +``` + ## Official Checkpoint Improvements The published weights are licensed under CC BY-ND 4.0. Modified weights, including adapters, deltas, pruned variants, and quantized variants, may not be redistributed. diff --git a/README.md b/README.md index e6e6de6..75be048 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. | 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..95f920a 100644 --- a/biofoundation/core/__init__.py +++ b/biofoundation/core/__init__.py @@ -20,11 +20,17 @@ """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, split_state_dict_by_prefix +from biofoundation.core.protocols import PredictionHead, SignalEncoder __all__ = [ "BatchRequirements", + "PredictionHead", + "SafetensorsCheckpointMixin", "SignalBatch", + "SignalEncoder", "as_signal_batch", "require_batch_fields", + "split_state_dict_by_prefix", ] diff --git a/biofoundation/core/batch.py b/biofoundation/core/batch.py index 5f44502..0fbd060 100644 --- a/biofoundation/core/batch.py +++ b/biofoundation/core/batch.py @@ -29,24 +29,54 @@ 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 + channel_coords: bool = False sensor_type: bool = False + num_padded_channels: bool = False + num_padded_timesteps: bool = False def as_signal_batch(batch: Any) -> SignalBatch: @@ -80,7 +110,10 @@ def require_batch_fields(batch: SignalBatch, requirements: BatchRequirements) -> ("label", requirements.label), ("channel_names", requirements.channel_names), ("channel_locations", requirements.channel_locations), + ("channel_coords", requirements.channel_coords), ("sensor_type", requirements.sensor_type), + ("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..c22e9ee --- /dev/null +++ b/biofoundation/core/checkpoints.py @@ -0,0 +1,113 @@ +#*----------------------------------------------------------------------------* +#* 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`` loads pre-trained weights by calling ``load_pretrained_checkpoint`` +for a Lightning ``.ckpt`` and ``load_safetensors_checkpoint`` for a ``.safetensors`` +file. Tasks that implement Lightning's own ``load_from_checkpoint`` instead can mix +this class in to expose both names without duplicating either loader. + +The mixin is deliberately thin: it maps names and formats, and delegates the actual +tensor matching to the task. Loading policy, including which shape mismatches are +tolerated, stays with the task that owns the model. +""" + +from typing import Any, Dict, Optional, Protocol + + +class _CheckpointLoadable(Protocol): # pragma: no cover - typing only + """The single method a task must provide for the mixin to delegate to.""" + + def load_from_checkpoint(self, checkpoint_path: str, **kwargs: Any) -> Any: ... + + +class SafetensorsCheckpointMixin: + """Expose BioFoundation's two checkpoint entry points over ``load_from_checkpoint``. + + Mix in ahead of :class:`pytorch_lightning.LightningModule` on tasks whose loading + logic already lives in ``load_from_checkpoint``: + + .. code-block:: python + + class MyTask(SafetensorsCheckpointMixin, pl.LightningModule): + def load_from_checkpoint(self, checkpoint_path, **kwargs): ... + + ``load_safetensors_checkpoint`` converts the flat ``.safetensors`` mapping into the + ``{"state_dict": ...}`` layout ``load_from_checkpoint`` expects, writes it to a + temporary file, and delegates. Keys are given a ``model.`` prefix when they lack + one, 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_from_checkpoint``.""" + + return self.load_from_checkpoint(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_from_checkpoint(checkpoint_path=str(converted), **kwargs) + + +def split_state_dict_by_prefix( + state_dict: Dict[str, Any], + prefixes: tuple[str, ...] = ("model_head.", "model."), +) -> Dict[str, Dict[str, Any]]: + """Group a flat Lightning ``state_dict`` by top-level module prefix. + + Args: + state_dict: Mapping of parameter name to tensor, as stored by Lightning. + prefixes: Prefixes to split on, longest first so that ``model_head.`` is + matched before ``model.``. + + Returns: + Mapping of prefix-without-the-dot to a prefix-stripped state dict. Keys that + match no prefix are collected under ``""``. + """ + + grouped: Dict[str, Dict[str, Any]] = {prefix.rstrip("."): {} for prefix in prefixes} + grouped[""] = {} + + for key, value in state_dict.items(): + for prefix in prefixes: + if key.startswith(prefix): + grouped[prefix.rstrip(".")][key[len(prefix):]] = value + break + else: + grouped[""][key] = value + + return grouped + + +__all__ = ["SafetensorsCheckpointMixin", "split_state_dict_by_prefix"] diff --git a/biofoundation/core/protocols.py b/biofoundation/core/protocols.py new file mode 100644 index 0000000..489feb9 --- /dev/null +++ b/biofoundation/core/protocols.py @@ -0,0 +1,81 @@ +#*----------------------------------------------------------------------------* +#* 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. + +BioFoundation carries two model shapes. The five original families bundle their +output layer into the model itself, selecting it from ``num_classes`` at +construction time. Newer families separate the two: an encoder that produces token +embeddings, and a prediction head that consumes them, each instantiated from its own +Hydra group. + +The protocols below describe the second shape. They are :class:`typing.Protocol` +definitions, so conformance is structural: a module satisfies one by having a +matching ``forward``, with no base class to inherit and no registration step. That +keeps the bundled families valid exactly as they are while giving the split families +a contract a type checker can verify. + +Tensors are annotated only under :data:`typing.TYPE_CHECKING`. This module is +imported by the fast contract test suite, which runs without PyTorch installed, so +nothing here may import ``torch`` at runtime. +""" + +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", *args: Any, **kwargs: Any) -> "Tensor": + """Encode a patched biosignal into token embeddings.""" + ... + + +@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", *args: Any, **kwargs: Any) -> "Tensor": + """Map token embeddings to a prediction.""" + ... + + +__all__ = ["PredictionHead", "SignalEncoder"] diff --git a/biofoundation/model_registry.py b/biofoundation/model_registry.py index 7b801f2..c37bcc0 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", "small", "base"), ), } ) 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..6e071ff --- /dev/null +++ b/config/data_module/pretrain_data_module_SCEReBrO.yaml @@ -0,0 +1,114 @@ +# @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.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: '${env:DATA_PATH}/pretraining/TUEG/TUEG.lmdb' + keys_path: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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: '${env:DATA_PATH}/pretraining/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..21af2b9 --- /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. 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/Mumtaz +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/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..200b650 --- /dev/null +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -0,0 +1,126 @@ +# @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 is supplied by the selected config/model group, so it always matches +# the encoder actually being built. + +# Where snapshot_download placed the Hugging Face release. Kept separate from the +# checkpoint path so a run can point at one file without restating the root: +# pretrained_safetensors_path='${pretrained_root}/SCEReBrO_${model_size}.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 + +# The corpus is selected as a single config group. Each file in config/dataset/ owns +# everything that varies per corpus: its path, whether samples are windows or +# sequences, whether labels are classes or a continuous target, the channel count, and +# the matching prediction head, task and criterion. Switch corpus with one override: +# python -u run_train.py +experiment=SCEReBrO_finetune 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 + +# Normalisation for S-CEReBrO is min-max to [-1, 1], applied per channel by +# LMDBDataset via apply_minmax and written into TUEG offline by make_tueg. The +# quantile normalisation that config/defaults.yaml configures belongs to the tasks the +# bundled families use; the S-CEReBrO tasks never read input_normalization. It is +# disabled here so the resolved config cannot be misread as enabling it. +input_normalization: + normalize: False + +defaults: + # The dataset group brings its own model_head, task and criterion, so this experiment + # deliberately does not select them. Hydra applies a config's own values after its + # defaults list, so anything set here would win over the group and the two could + # disagree; leaving them to config/dataset/ makes one file the single owner. + # + # model_head is not in the global defaults, so the dataset group adds it rather than + # overriding it, which leaves the five bundled families composing exactly as before. + # Hydra requires additions before overrides in a defaults list. + - /dataset: tuab + - override /data_module: finetune_data_module_SCEReBrO + - override /model: SCEReBrO_tiny # or SCEReBrO_small, SCEReBrO_base + - 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..3afc3fc --- /dev/null +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -0,0 +1,109 @@ +# @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 is supplied by the selected config/model group, so it always matches +# the encoder actually being built. + +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 + +# Normalisation for S-CEReBrO is min-max to [-1, 1], applied per channel by +# LMDBDataset via apply_minmax and written into TUEG offline by make_tueg. The +# quantile normalisation that config/defaults.yaml configures belongs to the tasks the +# bundled families use; the S-CEReBrO tasks never read input_normalization. It is +# disabled here so the resolved config cannot be misread as enabling it. +input_normalization: + normalize: False + +defaults: + # model_head is not part of the global defaults, so it is added here rather than + # overridden, leaving the composition of the five bundled families untouched. + # Hydra requires additions before overrides in a defaults list. + - /model_head: patch_reconstruction_head + - override /data_module: pretrain_data_module_SCEReBrO + - override /model: SCEReBrO_tiny # or SCEReBrO_small, SCEReBrO_base + - 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_base.yaml b/config/model/SCEReBrO_base.yaml new file mode 100644 index 0000000..c570d64 --- /dev/null +++ b/config/model/SCEReBrO_base.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 * +#*----------------------------------------------------------------------------* +# Declared here rather than in the experiment so that selecting this group with +# model=SCEReBrO_base also updates model_size, which names the output directory +# and the pre-trained checkpoint file. The two cannot drift apart. +model_size: base + +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: 400 + depth: 12 + num_heads: 16 + 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/SCEReBrO_small.yaml b/config/model/SCEReBrO_small.yaml new file mode 100644 index 0000000..3adacfb --- /dev/null +++ b/config/model/SCEReBrO_small.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 * +#*----------------------------------------------------------------------------* +# Declared here rather than in the experiment so that selecting this group with +# model=SCEReBrO_small also updates model_size, which names the output directory +# and the pre-trained checkpoint file. The two cannot drift apart. +model_size: small + +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: 200 + depth: 12 + num_heads: 10 + 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/SCEReBrO_tiny.yaml b/config/model/SCEReBrO_tiny.yaml new file mode 100644 index 0000000..26b7757 --- /dev/null +++ b/config/model/SCEReBrO_tiny.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 * +#*----------------------------------------------------------------------------* +# Declared here rather than in the experiment so that selecting this group with +# model=SCEReBrO_tiny also updates model_size, which names the output directory +# and the pre-trained checkpoint file. The two cannot drift apart. +model_size: tiny + +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..6b996f7 --- /dev/null +++ b/criterion/ce_criterion.py @@ -0,0 +1,55 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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..3605c20 --- /dev/null +++ b/criterion/focal_criterion.py @@ -0,0 +1,63 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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. + """ + 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..af85a1e --- /dev/null +++ b/criterion/masked_reconstruction_loss.py @@ -0,0 +1,82 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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 = per_patch[token_mask & attn_mask].mean() + logs = {"masked_loss": masked.item()} + + if self.alpha == 0: + return masked, logs + + visible = per_patch[(~token_mask) & attn_mask].mean() + logs["visible_loss"] = visible.item() + return masked + self.alpha * visible, logs diff --git a/criterion/mse_criterion.py b/criterion/mse_criterion.py new file mode 100644 index 0000000..9cfe9ca --- /dev/null +++ b/criterion/mse_criterion.py @@ -0,0 +1,46 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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..2b3ffe0 --- /dev/null +++ b/data_module/finetuning_data_module.py @@ -0,0 +1,84 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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..5ede1b2 --- /dev/null +++ b/data_module/pretraining_data_module.py @@ -0,0 +1,98 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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..7931f69 --- /dev/null +++ b/datasets/lmdb_dataset.py @@ -0,0 +1,276 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +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 + + +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: + """Open the LMDB environment lazily, on first read. + + Construction deliberately leaves ``env`` unset. LMDB refuses to open the same + path twice in one process, and ``run_train.py`` re-instantiates the data module + before the rank-zero test pass while the original may still be referenced. It + also avoids handing an open handle to forked dataloader workers. This mirrors + how :class:`~datasets.tueg_dataset.TUEGDataset` manages its environment. + """ + if self.env is None: + self.env = lmdb.open( + self.lmdb_path, readonly=True, lock=False, readahead=True, + max_readers=64, map_async=True, + ) + + 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] + + # No index file: scan once through a temporary environment, then close it so + # construction still leaves no handle open. + env = lmdb.open(self.lmdb_path, readonly=True, lock=False, readahead=True) + try: + with env.begin(buffers=True) as txn: + return [bytes(key) for key in txn.cursor().iternext(keys=True, values=False)] + finally: + env.close() + + 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 + 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): + """Close the LMDB environment.""" + env = getattr(self, "env", None) + if env is not None: + try: + env.close() + except Exception: + pass diff --git a/datasets/tueg_dataset.py b/datasets/tueg_dataset.py new file mode 100644 index 0000000..138c86f --- /dev/null +++ b/datasets/tueg_dataset.py @@ -0,0 +1,146 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import os +from collections import OrderedDict +from typing import Any, Dict, List, Optional + +import lmdb +import numpy as np +import torch +from torch.utils.data import Dataset + + +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: + """Open the LMDB environment lazily so it is not inherited across worker forks.""" + if self.env is None: + self.env = lmdb.open( + self.lmdb_path, readonly=True, lock=False, readahead=True, map_async=True + ) + + 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): + """Close the LMDB environment.""" + env = getattr(self, "env", None) + if env is not None: + try: + env.close() + except Exception: + pass diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md index 630f1ef..ebe11ac 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}: Windowed Alternating Attention for Compact {EEG} Representation Learning}, + author={Bucagu, Glenn Anta and Dimofte, Alexandru 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..b3e98f0 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,6 +40,7 @@ The training guide covers environment variables, Hydra experiment selection, the ### 5. Project References - [`CONTRIBUTING.md`](../CONTRIBUTING.md) defines extension and pull request expectations. +- [`docs/adr`](./adr/) records architecture decisions affecting the shared contracts. - [`CITATIONS.md`](./CITATIONS.md) contains BibTeX for all five 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..05f6393 --- /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. + +### 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 follows the BESA electrode and surface location tables and assigns fixed coordinates to reference electrodes that have no scalp position of their own. + +```bash +python -m make_datasets.make_tuab --output $DATA_PATH/finetuning/TUAB +python -m make_datasets.make_tueg --output $DATA_PATH/pretraining/TUEG +``` + +### Architecture Overview + +| Stage | Module | +| --- | --- | +| Tokenisation | `TemporalConvTokenizer`: three strided 1D convolutions per `(channel, patch)` pair, projected to `embed_dim` | +| Position | Learned temporal table, shared across channels, sliced to the patches actually present | +| Channel | Shared MLP over each electrode's 3D coordinate, halves concatenated | +| Backbone | `depth` pre-norm transformer blocks with alternating attention | +| Output | A separate `PredictionHead` | + +Attention alternates by block index: + +| Block | Attends over | Window | +| --- | --- | --- | +| even | channels, at a fixed patch position (spatial) | `window_size_spatial`, dilated by `dilation_cycle_spatial`, shifted by `shift_cycle_spatial` | +| odd | patch positions, within a fixed channel (temporal) | `window_size_temporal`, dilated by `dilation_cycle_temporal`, shifted by `shift_cycle_temporal` | + +Dilation and shift schedules are indexed by spatial/temporal *pair*, so a spatial block and the temporal block after it share a schedule entry. Setting `use_axial_mode: True` runs all spatial blocks before all temporal blocks instead. + +`attention_type` selects the mechanism: `windowed-alternating` is the published method; `alternating` (no windowing) and `full` (all tokens at once) are the ablation baselines. + +Padded channels are replaced by a learned pad token, excluded from masking, and masked out of attention, so montages of different sizes share one batch safely. + +### Self-Supervised Learning (SSL) Objective + +SimMIM-style masked reconstruction. Waveforms are patched and embedded, a random subset of real tokens is replaced by a learned mask token, and the encoder sees the full sequence of visible and masked tokens in their original order. A linear decoder reconstructs every patch and the loss is taken over the masked, non-padded positions. `alpha` adds a weighted term over visible patches, which stabilises early training. + +### Downstream Tasks + +| Layout | Head | Used by | +| --- | --- | --- | +| Window classification | `MlpClassificationHead` | TUAB, CHB-MIT, Neonate, PhysioNet-MI, SHU-MI, STEW, Mumtaz, MentalArithmetic, SEED-V | +| Sequence classification | `SequenceClassificationHead` | ISRUC sleep staging | +| Scalar regression | `MlpRegressionHead` | SEED-VIG vigilance | + +Prepared datasets and their shapes: + +| Dataset | Task | Channels | Window | Classes | +| --- | --- | --- | --- | --- | +| TUAB | binary classification | 22 | 10 s | 2 | +| CHB-MIT | seizure detection | 16 | 10 s | 2 | +| Neonate | seizure detection | 18 | 5 s | 2 | +| PhysioNet-MI | motor imagery | 64 | 4 s | 4 | +| SHU-MI | motor imagery | 32 | 4 s | 2 | +| STEW | workload | 14 | 4 s | 3 | +| Mumtaz | depression | 20 | 5 s | 2 | +| MentalArithmetic | mental arithmetic | 20 | 5 s | 2 | +| SEED-V | emotion | 62 | 4 s | 5 | +| ISRUC | sleep staging | 6 | 30 s x 20 epochs | 5 | +| SEED-VIG | vigilance regression | 17 | 8 s | continuous | + +### Model Variants + +| Variant | `embed_dim` | `depth` | `num_heads` | Config | +| --- | --- | --- | --- | --- | +| tiny | 180 | 6 | 5 | [`SCEReBrO_tiny`](../../config/model/SCEReBrO_tiny.yaml) | +| small | 200 | 12 | 10 | [`SCEReBrO_small`](../../config/model/SCEReBrO_small.yaml) | +| base | 400 | 12 | 16 | [`SCEReBrO_base`](../../config/model/SCEReBrO_base.yaml) | + +### Training Setup + +Pre-train on the union of the prepared corpora: + +```bash +python -u run_train.py +experiment=SCEReBrO_pretrain +``` + +Fine-tune on TUAB, the default corpus: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune \ + pretrained_safetensors_path=/absolute/path/to/SCEReBrO_tiny.safetensors +``` + +Fine-tune on another corpus by selecting it. Each file in [`config/dataset`](../../config/dataset/) owns everything that varies per corpus: its path, whether samples are windows or sequences, whether labels are classes or a continuous target, the channel count, and the matching prediction head, task and criterion. + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune dataset=chb-mit +python -u run_train.py +experiment=SCEReBrO_finetune dataset=isruc +python -u run_train.py +experiment=SCEReBrO_finetune dataset=seed-vig +``` + +| `dataset=` | Task | Channels | Window | Classes | Head | +| --- | --- | --- | --- | --- | --- | +| `tuab` | binary classification | 22 | 10 s | 2 | `MlpClassificationHead` | +| `chb-mit` | seizure detection | 16 | 10 s | 2 | `MlpClassificationHead` | +| `neonate` | seizure detection | 18 | 5 s | 2 | `MlpClassificationHead` | +| `physionet-mi` | motor imagery | 64 | 4 s | 4 | `MlpClassificationHead` | +| `shu-mi` | motor imagery | 32 | 4 s | 2 | `MlpClassificationHead` | +| `stew` | workload | 14 | 4 s | 3 | `MlpClassificationHead` | +| `mumtaz` | depression | 20 | 5 s | 2 | `MlpClassificationHead` | +| `mental-arithmetic` | mental arithmetic | 20 | 5 s | 2 | `MlpClassificationHead` | +| `seed-v` | emotion | 62 | 4 s | 5 | `MlpClassificationHead` | +| `isruc` | sleep staging | 6 | 30 s x 20 epochs | 5 | `SequenceClassificationHead` | +| `seed-vig` | vigilance regression | 17 | 8 s | continuous | `MlpRegressionHead` | + +Individual values can still be overridden on top of a selection, for a montage that differs from the prepared one: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune dataset=tuab model.num_channels=19 +``` + +**Where each setting lives.** The channel count is `model.num_channels`, and it must equal what the dataset yields; the encoder raises rather than reshaping if they disagree. The window length is not configured at all for window classification with the default mean pooling, because the task derives the patch count from the data and the encoder slices its position table to match; the only limit is `max_timesteps / patch_size`, which is 30 s, and exceeding it raises. `model_head.num_patches` is read only by `SequenceClassificationHead` and by `MlpClassificationHead` when `pooling_method` is `flatten`. The task type is not a separate flag: selecting a dataset selects the head, task and criterion together. + +**Adding a corpus.** Copy the closest file in `config/dataset/`, set its path, channel count and label details, and leave the head, task and criterion selections alone unless the task type differs. Do not set these values in the experiment: Hydra applies a config's own values after its defaults list, so a key set in both places resolves to the experiment's copy and the dataset file is silently ignored. A contract test enforces this. + +A linear-probe style run freezes the encoder blocks while leaving tokenisation and the embeddings trainable: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune task.freeze_backbone=True +``` + +Fine-tuning uses layer-wise learning-rate decay: blocks closer to the input receive `lr * decay ** (depth - 1 - block_idx)`. Biases, normalisation weights, and the embedding tables are excluded from weight decay, and the head forms its own parameter group. + +### Smoke Test With Synthetic Data + +To check the pipeline end to end without prepared data, generate synthetic corpora in +the exact on-disk formats the readers expect. The signals are band-limited noise and +the labels are random, so this verifies that a run works, not that it learns anything. + +```bash +export DATA_PATH=/absolute/path/to/dummy-data +export CHECKPOINT_DIR=/absolute/path/to/experiments + +python -m make_datasets.make_dummy_scerebro_dataset --output $DATA_PATH \ + --pretrain-samples 24 --finetune-samples 32 +``` + +Add `--datasets pretrain tuab isruc seed-vig` to also generate the sequence and +regression corpora. + +Pre-train, then fine-tune from the resulting checkpoint: + +```bash +python -u run_train.py +experiment=SCEReBrO_pretrain \ + trainer.accelerator=cpu trainer.devices=1 trainer.strategy=auto \ + trainer.max_epochs=1 trainer.accumulate_grad_batches=1 \ + trainer.check_val_every_n_epoch=1 scheduler.warmup_epochs=0 \ + batch_size=4 num_workers=0 final_validate=False + +python -u run_train.py +experiment=SCEReBrO_finetune \ + pretrained_checkpoint_path=$CHECKPOINT_DIR/checkpoints/SCEReBrO_pretrain//last.ckpt \ + trainer.accelerator=cpu trainer.devices=1 trainer.strategy=auto \ + trainer.max_epochs=2 scheduler.warmup_epochs=0 \ + batch_size=4 num_workers=0 +``` + +Drop the `trainer.*` and `num_workers` overrides on a GPU machine; they exist only to +make the run finish quickly on CPU. + +On PyTorch 2.6 and newer, reloading a checkpoint for the final validation and test +passes fails with `UnpicklingError: Weights only load failed`, because `torch.load` +now defaults to `weights_only=True` and every task stores its Hydra configuration in +the checkpoint. This affects `run_train.py` for all model families, not only this one. +Either export `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` or pass `final_validate=False +final_test=False`. + +### Pretrained Weights + +The [PulpBio/S-CEReBrO Hugging Face repository](https://huggingface.co/PulpBio/S-CEReBrO) provides tiny, small and base checkpoints matching the model configs in [`config/model`](../../config/model/). The weights are licensed under CC BY-ND 4.0. + +`snapshot_download` writes to whatever `local_dir` you give it, resolved relative to the +current working directory when it is not absolute. The fine-tuning experiment expects +the release under `$CHECKPOINT_DIR/pretrained/S-CEReBrO`, which is the default value of +`pretrained_root`, so download it there: + +```python +import os +from huggingface_hub import snapshot_download + +snapshot_download( + repo_id="PulpBio/S-CEReBrO", + local_dir=os.path.join(os.environ["CHECKPOINT_DIR"], "pretrained", "S-CEReBrO"), +) +``` + +The checkpoint path can then be written as an interpolation instead of an absolute +path. `model_size` is declared by the selected `config/model` group, so it always +matches the encoder being built and the two cannot drift: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune model=SCEReBrO_tiny \ + 'pretrained_safetensors_path=${pretrained_root}/SCEReBrO_${model_size}.safetensors' +``` + +Switching size needs one change, and the checkpoint follows: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune model=SCEReBrO_base \ + 'pretrained_safetensors_path=${pretrained_root}/SCEReBrO_${model_size}.safetensors' +``` + +Single-quote the override so the shell does not expand `${...}` before Hydra sees it. +`pretrained_root` defaults to `${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO` and can be +pointed elsewhere with `pretrained_root=/some/other/dir`. An absolute +`pretrained_safetensors_path` still works exactly as it does for the other families. + +The size flag and the checkpoint must agree: loading a `base` checkpoint into a `tiny` +encoder is not an error, because shape-mismatched tensors are skipped rather than +forced. Check the `[load:model] loaded=... shape_mismatch=... unexpected=...` line the +loader prints; on a correct pairing `shape_mismatch` and `unexpected` are both zero and +`loaded` equals `total_target`. + +A Lightning `.ckpt` produced by a local pre-training run is passed with +`pretrained_checkpoint_path` instead: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune \ + pretrained_checkpoint_path=$CHECKPOINT_DIR/checkpoints/SCEReBrO_pretrain//last.ckpt +``` + +Convert one to safetensors for distribution with the repository's own tool: + +```bash +python util/ckpt_to_safetensor.py \ + --ckpt_path $CHECKPOINT_DIR/checkpoints/SCEReBrO_pretrain//last.ckpt \ + --safetensor_path SCEReBrO_tiny.safetensors +``` + +Only the encoder needs to transfer. Head weights are loaded from a checkpoint only when +`include_head=True` is requested, so a reconstruction head from pre-training never +overwrites a freshly initialised classification head. diff --git a/make_datasets/common.py b/make_datasets/common.py new file mode 100644 index 0000000..c215c64 --- /dev/null +++ b/make_datasets/common.py @@ -0,0 +1,325 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import argparse +import os +import pickle +from collections import Counter +from typing import Any, Callable, Dict, Iterable, List, Sequence, Tuple + +import lmdb +import numpy as np +from tqdm import tqdm + +from make_datasets import electrode_positions + +SAMPLING_FREQ = 200 +SPLITS = ("train", "val", "test") +DEFAULT_MAP_SIZE = 60 * 10 ** 9 +COMMIT_INTERVAL = 4000 + + +def electrode_coordinate(name: str) -> Tuple[float, float, float]: + """Return the 3D coordinate of one electrode, or the origin if it is unknown.""" + if name in electrode_positions.ELECTRODE_ANGLES: + angles = electrode_positions.ELECTRODE_ANGLES[name] + return electrode_positions.get_electrode_3d_positions(angles["theta"], angles["phi"]) + return electrode_positions.SPECIAL_REFERENCE_POSITIONS.get(name, (0.0, 0.0, 0.0)) + + +def referential_coordinates(channel_names: Sequence[str], reference: str) -> np.ndarray: + """Build channel coordinates for a montage where every channel shares one reference. + + Args: + channel_names: Electrode names in the order the channels appear in the data. + reference: Name of the shared reference electrode. + + Returns: + Array of shape ``(num_channels, 2, 3)`` holding the active and reference + coordinates of each channel. + """ + reference_position = electrode_coordinate(reference) + coords = np.zeros((len(channel_names), 2, 3), dtype=np.float32) + for index, name in enumerate(channel_names): + coords[index, 0, :] = electrode_coordinate(name) + coords[index, 1, :] = reference_position + return coords + + +def bipolar_coordinates(pairs: Sequence[Tuple[str, str]]) -> np.ndarray: + """Build channel coordinates for a bipolar montage. + + Args: + pairs: ``(active, reference)`` electrode names for each channel. + + Returns: + Array of shape ``(num_channels, 2, 3)``. + """ + coords = np.zeros((len(pairs), 2, 3), dtype=np.float32) + for index, (active, reference) in enumerate(pairs): + coords[index, 0, :] = electrode_coordinate(active) + coords[index, 1, :] = electrode_coordinate(reference) + return coords + + +def index_splits(items: Sequence[Any], train_end: int, val_end: int) -> Dict[str, List[Any]]: + """Split an ordered sequence of recordings into train, validation and test. + + Splitting on recordings rather than windows keeps every window of a recording, + and therefore of a subject, inside a single split. + + Args: + items: Recordings in a deterministic order. + train_end: Number of leading recordings used for training. + val_end: Index at which validation ends and test begins. + + Returns: + Mapping from split name to the recordings assigned to it. + """ + if not 0 < train_end <= val_end <= len(items): + raise ValueError( + f"Invalid split boundaries train_end={train_end} val_end={val_end} for {len(items)} recordings" + ) + return {"train": list(items[:train_end]), "val": list(items[train_end:val_end]), "test": list(items[val_end:])} + + +class LMDBWriter: + """Writes pickled EEG samples into one LMDB per split and reports the result. + + Transactions are committed every ``COMMIT_INTERVAL`` samples so that a long job does + not hold one unbounded write transaction. + + Args: + output_dir: Directory to create the LMDB files in. + splits: Split names to open an LMDB for. A single-element sequence produces one + pooled database, which is what the pre-training corpora use. + map_size: Maximum size of each LMDB in bytes. + dry_run: Collect statistics without writing anything. + """ + + def __init__( + self, + output_dir: str, + splits: Sequence[str] = SPLITS, + map_size: int = DEFAULT_MAP_SIZE, + dry_run: bool = False, + ): + self.output_dir = output_dir + self.splits = tuple(splits) + self.dry_run = dry_run + self.written = 0 + self.counts: Dict[str, int] = {split: 0 for split in self.splits} + self.labels: Dict[str, Counter] = {split: Counter() for split in self.splits} + self.shapes: Counter = Counter() + self.subjects: Dict[str, set] = {split: set() for split in self.splits} + + self.envs: Dict[str, Any] = {} + self.txns: Dict[str, Any] = {} + if not dry_run: + os.makedirs(output_dir, exist_ok=True) + for split in self.splits: + self.envs[split] = lmdb.open(os.path.join(output_dir, f"{split}.lmdb"), map_size=map_size) + self.txns[split] = self.envs[split].begin(write=True) + + def put(self, split: str, key: bytes, sample: Dict[str, Any]) -> None: + """Record one sample, writing it unless this is a dry run.""" + self.counts[split] += 1 + eeg = sample["eeg"] + self.shapes[tuple(eeg.shape)] += 1 + + label = sample.get("label") + if label is not None: + for value in np.atleast_1d(np.asarray(label)).ravel(): + self.labels[split][value.item()] += 1 + if "subject_id" in sample: + self.subjects[split].add(str(sample["subject_id"])) + + if self.dry_run: + return + + self.txns[split].put(key, pickle.dumps(sample)) + self.written += 1 + if self.written % COMMIT_INTERVAL == 0: + for name in self.splits: + self.txns[name].commit() + self.txns[name] = self.envs[name].begin(write=True) + + def close(self) -> None: + """Commit and close every LMDB.""" + if self.dry_run: + return + for split in self.splits: + self.txns[split].commit() + self.envs[split].close() + + def summarise(self, name: str) -> None: + """Print per-split sample counts, label distributions and window shapes.""" + verb = "computed" if self.dry_run else "written" + print(f"\n{name}: {sum(self.counts.values())} samples {verb} to {self.output_dir}") + for split in self.splits: + subjects = self.subjects[split] + suffix = f", {len(subjects)} subjects" if subjects else "" + print(f" {split:5s}: {self.counts[split]} samples{suffix}") + if self.labels[split]: + distribution = ", ".join( + f"{label}:{count}" for label, count in sorted(self.labels[split].items()) + ) + print(f" labels {distribution}") + print(" window shapes: " + ", ".join(f"{shape}:{count}" for shape, count in sorted(self.shapes.items()))) + if len(self.splits) > 1: + overlap = set.intersection(*(self.subjects[s] for s in self.splits)) if all( + self.subjects[s] for s in self.splits + ) else set() + if overlap: + print(f" WARNING: {len(overlap)} subjects appear in more than one split") + + +class PackedLMDBWriter: + """Writes fixed-size raw byte records into one LMDB alongside a key list. + + Used for the largest pre-training corpus, where pickling every sample is a + measurable overhead. Each value is a waveform followed by its channel + coordinates, both float32, so readers can slice the blob without unpickling. + + Args: + lmdb_path: Path of the LMDB to create. + keys_path: Path of the newline-separated key list to write. + map_size: Maximum size of the LMDB in bytes. + dry_run: Collect statistics without writing anything. + """ + + def __init__(self, lmdb_path: str, keys_path: str, map_size: int = DEFAULT_MAP_SIZE, dry_run: bool = False): + self.lmdb_path = lmdb_path + self.keys_path = keys_path + self.dry_run = dry_run + self.keys: List[str] = [] + self.written = 0 + self.env = None + self.txn = None + if not dry_run: + os.makedirs(os.path.dirname(lmdb_path) or ".", exist_ok=True) + self.env = lmdb.open(lmdb_path, map_size=map_size) + self.txn = self.env.begin(write=True) + + def put(self, key: str, waveform: np.ndarray, channel_coords: np.ndarray) -> None: + """Record one window as a packed byte blob.""" + self.keys.append(key) + if self.dry_run: + return + payload = waveform.astype(np.float32).tobytes() + channel_coords.astype(np.float32).tobytes() + self.txn.put(key.encode("ascii"), payload) + self.written += 1 + if self.written % COMMIT_INTERVAL == 0: + self.txn.commit() + self.txn = self.env.begin(write=True) + + def close(self) -> None: + """Commit the LMDB and write the key list.""" + if self.dry_run: + return + self.txn.commit() + self.env.close() + with open(self.keys_path, "w") as handle: + handle.write("\n".join(self.keys) + "\n") + + def summarise(self, name: str) -> None: + """Print the number of windows written.""" + verb = "computed" if self.dry_run else "written" + print(f"\n{name}: {len(self.keys)} windows {verb} to {self.lmdb_path}") + + +def build_arg_parser(description: str) -> argparse.ArgumentParser: + """Create the argument parser shared by every preprocessing script.""" + parser = argparse.ArgumentParser(description=description) + parser.add_argument("--input_dir", required=True, help="Directory holding the raw dataset") + parser.add_argument("--output_dir", required=True, help="Directory to write LMDB files into") + parser.add_argument("--num_workers", type=int, default=8, help="Parallel worker processes") + parser.add_argument("--dry_run", action="store_true", help="Report statistics without writing") + return parser + + +def run_jobs( + worker: Callable[[Any], Iterable[Any]], tasks: Sequence[Any], num_workers: int, description: str +) -> Iterable[Any]: + """Map ``worker`` over ``tasks``, yielding each result as it completes. + + Runs in the calling process when ``num_workers`` is 1, which keeps tracebacks + readable while debugging a new dataset. + """ + if num_workers <= 1: + for task in tqdm(tasks, desc=description): + yield worker(task) + return + + from multiprocessing import Pool + + with Pool(num_workers) as pool: + for result in tqdm(pool.imap_unordered(worker, tasks), total=len(tasks), desc=description): + yield result + + +def list_files(root: str, extensions: Sequence[str]) -> List[str]: + """Return every file under ``root`` with one of ``extensions``, sorted by relative path.""" + wanted = tuple(extension.lower() for extension in extensions) + found = [] + for directory, _, filenames in os.walk(root): + for filename in filenames: + if filename.lower().endswith(wanted): + found.append(os.path.relpath(os.path.join(directory, filename), root)) + return sorted(found) + + +def resample_to_target(signal_array: np.ndarray, source_freq: int, axis: int = -1) -> np.ndarray: + """Resample along ``axis`` from ``source_freq`` to the project-wide 200 Hz.""" + if source_freq == SAMPLING_FREQ: + return signal_array + from scipy import signal as scipy_signal + + num_samples = int(round(signal_array.shape[axis] * SAMPLING_FREQ / source_freq)) + return scipy_signal.resample(signal_array, num_samples, axis=axis) + +def slice_windows(data: np.ndarray, window_samples: int) -> np.ndarray: + """Cut a ``(channels, timesteps)`` recording into non-overlapping windows. + + Returns: + Array of shape ``(num_windows, channels, window_samples)``; empty when the + recording is shorter than one window. + """ + channels, timesteps = data.shape + num_windows = timesteps // window_samples + if num_windows == 0: + return np.empty((0, channels, window_samples), dtype=np.float32) + trimmed = data[:, : num_windows * window_samples] + return trimmed.reshape(channels, num_windows, window_samples).transpose(1, 0, 2).astype(np.float32) + + +def write_pretraining_corpus(name, tasks, worker, output_dir, num_workers, dry_run): + """Run a pre-training corpus job and write every window into one pooled LMDB. + + Pre-training corpora carry no labels and no splits: the pre-training data module + holds out a fraction of windows for validation itself. + """ + writer = LMDBWriter(output_dir, splits=("all",), dry_run=dry_run) + for samples in run_jobs(worker, tasks, num_workers, f"{name} recordings"): + for key, sample in samples: + writer.put("all", key, sample) + writer.close() + writer.summarise(name) diff --git a/make_datasets/electrode_positions.py b/make_datasets/electrode_positions.py new file mode 100644 index 0000000..dde4ba8 --- /dev/null +++ b/make_datasets/electrode_positions.py @@ -0,0 +1,179 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""3D electrode coordinates for the standard 118-electrode layout. + +Angular positions follow the BESA electrode and surface location tables at +https://wiki.besa.de/index.php?title=Electrodes_and_Surface_Locations. Reference +electrodes that have no scalp position of their own are given fixed coordinates. +""" + +import numpy as np + + +SPECIAL_REFERENCE_POSITIONS = { + 'AR': (0.0, 0.0, 0.2), + 'LE': (0, 0, -0.615), +} + +ELECTRODE_ANGLES = { + 'F11': {'theta': -130, 'phi': -40}, + 'F9': {'theta': -115, 'phi': -35}, + 'F7': {'theta': -92, 'phi': -36}, + 'F5': {'theta': -75, 'phi': -41}, + 'F3': {'theta': -60, 'phi': -51}, + 'F1': {'theta': -50, 'phi': -68}, + 'FZ': {'theta': 46, 'phi': 90}, + 'F2': {'theta': 50, 'phi': 68}, + 'F4': {'theta': 60, 'phi': 51}, + 'F6': {'theta': 75, 'phi': 41}, + 'F8': {'theta': 92, 'phi': 36}, + 'F10': {'theta': 115, 'phi': 35}, + 'F12': {'theta': 130, 'phi': 40}, + 'FT11': {'theta': -130, 'phi': -22}, + 'FT9': {'theta': -115, 'phi': -18}, + 'FT7': {'theta': -92, 'phi': -18}, + 'FC5': {'theta': -71, 'phi': -21}, + 'FC3': {'theta': -50, 'phi': -28}, + 'FC1': {'theta': -32, 'phi': -45}, + 'FCZ': {'theta': 23, 'phi': 90}, + 'FC2': {'theta': 32, 'phi': 45}, + 'FC4': {'theta': 50, 'phi': 28}, + 'FC6': {'theta': 71, 'phi': 21}, + 'FT8': {'theta': 92, 'phi': 18}, + 'FT10': {'theta': 115, 'phi': 18}, + 'FT12': {'theta': 130, 'phi': 22}, + 'T9': {'theta': -115, 'phi': 0}, + 'LPA': {'theta': -115, 'phi': 0}, + 'T7': {'theta': -92, 'phi': 0}, + 'C5': {'theta': -69, 'phi': 0}, + 'C3': {'theta': -46, 'phi': 0}, + 'C1': {'theta': -23, 'phi': 0}, + 'CZ': {'theta': 0, 'phi': 0}, + 'C2': {'theta': 23, 'phi': 0}, + 'C4': {'theta': 46, 'phi': 0}, + 'C6': {'theta': 69, 'phi': 0}, + 'T8': {'theta': 92, 'phi': 0}, + 'T10': {'theta': 115, 'phi': 0}, + 'RPA': {'theta': 115, 'phi': 0}, + 'P11': {'theta': -130, 'phi': 40}, + 'P9': {'theta': -115, 'phi': 36}, + 'P7': {'theta': -92, 'phi': 36}, + 'P5': {'theta': -75, 'phi': 41}, + 'P3': {'theta': -60, 'phi': 51}, + 'P1': {'theta': -50, 'phi': 68}, + 'PZ': {'theta': 46, 'phi': -90}, + 'P2': {'theta': 50, 'phi': -68}, + 'P4': {'theta': 60, 'phi': -51}, + 'P6': {'theta': 75, 'phi': -41}, + 'P8': {'theta': 92, 'phi': -36}, + 'P10': {'theta': 115, 'phi': -36}, + 'P12': {'theta': 130, 'phi': -40}, + 'TP9': {'theta': -115, 'phi': 18}, + 'TP7': {'theta': -92, 'phi': 18}, + 'CP5': {'theta': -71, 'phi': 21}, + 'CP3': {'theta': -50, 'phi': 28}, + 'CP1': {'theta': -32, 'phi': 45}, + 'CPZ': {'theta': 23, 'phi': -90}, + 'CP2': {'theta': 32, 'phi': -45}, + 'CP4': {'theta': 50, 'phi': -28}, + 'CP6': {'theta': 71, 'phi': -21}, + 'TP8': {'theta': 92, 'phi': -18}, + 'TP10': {'theta': 115, 'phi': -18}, + 'AF9': {'theta': -115, 'phi': -47}, + 'AF7': {'theta': -92, 'phi': -52}, + 'AF5': {'theta': -83, 'phi': -59}, + 'AF3': {'theta': -74, 'phi': -67}, + 'AF1': {'theta': -71, 'phi': -78}, + 'AFZ': {'theta': 69, 'phi': 90}, + 'AF2': {'theta': 71, 'phi': 78}, + 'AF4': {'theta': 74, 'phi': 67}, + 'AF6': {'theta': 83, 'phi': 59}, + 'AF8': {'theta': 92, 'phi': 52}, + 'AF10': {'theta': 115, 'phi': 47}, + 'FP1': {'theta': -92, 'phi': -72}, + 'FPZ': {'theta': 92, 'phi': 90}, + 'FP2': {'theta': 92, 'phi': 72}, + 'NZ': {'theta': 112, 'phi': 90}, + 'NAS': {'theta': 112, 'phi': 90}, + 'Chin': {'theta': 155, 'phi': 90}, + 'LO1': {'theta': -118, 'phi': -48}, + 'LO2': {'theta': 118, 'phi': 48}, + 'SO1': {'theta': -105, 'phi': -65}, + 'SO2': {'theta': 105, 'phi': 65}, + 'IO1': {'theta': -125, 'phi': -63}, + 'IO2': {'theta': 125, 'phi': 63}, + 'T3': {'theta': -92, 'phi': 0}, + 'T4': {'theta': 92, 'phi': 0}, + 'T5': {'theta': -92, 'phi': 36}, + 'T6': {'theta': 92, 'phi': -36}, + 'A1': {'theta': -128, 'phi': 3}, + 'A2': {'theta': 128, 'phi': -3}, + 'T1': {'theta': -108, 'phi': -20}, + 'T2': {'theta': 108, 'phi': 20}, + 'O1': {'theta': -92, 'phi': 72}, + 'OZ': {'theta': 92, 'phi': -90}, + 'O2': {'theta': 92, 'phi': -72}, + 'O9': {'theta': -115, 'phi': 72}, + 'O10': {'theta': 115, 'phi': -72}, + 'CB1': {'theta': -130, 'phi': 45}, + 'CB2': {'theta': 130, 'phi': -45}, + 'IZ': {'theta': 115, 'phi': -90}, + 'Neck': {'theta': 150, 'phi': -90}, + 'SP1': {'theta': -145, 'phi': -25}, + 'SP2': {'theta': 145, 'phi': 25}, + 'M1': {'theta': -120, 'phi': 25}, + 'M2': {'theta': 120, 'phi': -25}, + 'PO9': {'theta': -115, 'phi': 54}, + 'PO7': {'theta': -92, 'phi': 54}, + 'PO5': {'theta': -83, 'phi': 59}, + 'PO3': {'theta': -74, 'phi': 67}, + 'PO1': {'theta': -71, 'phi': 78}, + 'POZ': {'theta': 69, 'phi': -90}, + 'PO2': {'theta': 71, 'phi': -78}, + 'PO4': {'theta': 74, 'phi': -67}, + 'PO6': {'theta': 83, 'phi': -59}, + 'PO8': {'theta': 92, 'phi': -54}, + 'PO10': {'theta': 115, 'phi': -54} +} + + +def get_electrode_3d_positions(theta, phi, radius=1): + """ + Converts spherical coordinates (theta, phi) to Cartesian coordinates (x, y, z). + We use radius = 1 to have to normalize each coordinate in the range [0, 1] + + Args: + theta (float): Azimuthal angle in degrees. + phi (float): Polar angle in degrees. + radius (float): Radius of the sphere. + + Returns: + tuple: (x, y, z) 3D Cartesian coordinates. + """ + theta_rad = np.radians(theta) + phi_rad = np.radians(phi) + + x = radius * np.cos(phi_rad) * np.cos(theta_rad) + y = radius * np.cos(phi_rad) * np.sin(theta_rad) + z = radius * np.sin(phi_rad) + + return x, y, z diff --git a/make_datasets/make_bci_ner.py b/make_datasets/make_bci_ner.py new file mode 100644 index 0000000..2cab2d5 --- /dev/null +++ b/make_datasets/make_bci_ner.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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the BCI Challenge NER pre-training corpus into a pooled LMDB.""" + +import os + +import mne +import pandas as pd + +from make_datasets.common import ( + SAMPLING_FREQ, + build_arg_parser, + electrode_coordinate, + list_files, + slice_windows, + write_pretraining_corpus, +) +import numpy as np + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS +EXCLUDED_COLUMNS = ["EOG", "FeedBackEvent"] +REFERENCE = "AR" + + +def channel_coordinates(channel_names): + """Build referential coordinates for the channels present in one session file.""" + reference_position = electrode_coordinate(REFERENCE) + coords = np.zeros((len(channel_names), 2, 3), dtype=np.float32) + for index, name in enumerate(channel_names): + coords[index, 0, :] = electrode_coordinate(name.strip().upper()) + coords[index, 1, :] = reference_position + return coords + + +def process_recording(task): + """Filter and slice one comma-separated session recording.""" + root, relative_path = task + frame = pd.read_csv(os.path.join(root, relative_path)) + if frame.shape[0] == 0: + return [] + + columns = [name for name in frame.columns[1:] if name not in EXCLUDED_COLUMNS] + data = frame[columns].to_numpy().T + + info = mne.create_info(ch_names=columns, sfreq=SAMPLING_FREQ, ch_types="eeg", verbose=False) + raw = mne.io.RawArray(data, info, verbose=False) + raw.filter(l_freq=0.5, h_freq=30, method="fir", picks="eeg", verbose=False) + + coords = channel_coordinates(columns) + stem = os.path.splitext(os.path.basename(relative_path))[0] + return [ + ( + f"{stem}-{index}".encode(), + {"eeg": window, "channel_coords": coords, "subject_id": stem}, + ) + for index, window in enumerate(slice_windows(raw.get_data(), WINDOW_SAMPLES)) + ] + + +def main(): + """Slice BCI-NER into 30-second pre-training windows.""" + args = build_arg_parser("BCI-NER pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + tasks = [(args.input_dir, name) for name in list_files(args.input_dir, [".csv"])] + write_pretraining_corpus( + "BCI-NER", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_boas.py b/make_datasets/make_boas.py new file mode 100644 index 0000000..ae21e56 --- /dev/null +++ b/make_datasets/make_boas.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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the BOAS headband pre-training corpus into a pooled LMDB.""" + +import os + +import mne + +from make_datasets.common import ( + SAMPLING_FREQ, + build_arg_parser, + list_files, + referential_coordinates, + slice_windows, + write_pretraining_corpus, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS + +CHANNEL_RENAMES = {"HB_1": "AF7", "HB_2": "AF8"} +CHANNELS = ["AF7", "AF8"] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + + +def process_recording(task): + """Rename the headband channels, filter and slice one recording.""" + root, relative_path = task + raw = mne.io.read_raw_edf(os.path.join(root, relative_path), preload=True, verbose=False) + raw.rename_channels({key: value for key, value in CHANNEL_RENAMES.items() if key in raw.ch_names}) + + present = [name for name in CHANNELS if name in raw.ch_names] + if len(present) != NUM_CHANNELS: + return [] + raw.pick(present) + raw.reorder_channels(CHANNELS) + raw.filter(l_freq=0.1, h_freq=10, method="fir", picks="eeg", verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, verbose=False) + + stem = os.path.splitext(os.path.basename(relative_path))[0] + return [ + ( + f"{stem}-{index}".encode(), + {"eeg": window, "channel_coords": CHANNEL_COORDS, "subject_id": stem}, + ) + for index, window in enumerate(slice_windows(raw.get_data(), WINDOW_SAMPLES)) + ] + + +def main(): + """Slice BOAS into 30-second pre-training windows.""" + args = build_arg_parser("BOAS pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + tasks = [(args.input_dir, name) for name in list_files(args.input_dir, [".edf"])] + write_pretraining_corpus( + "BOAS", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_chb_mit.py b/make_datasets/make_chb_mit.py new file mode 100644 index 0000000..a2dda6f --- /dev/null +++ b/make_datasets/make_chb_mit.py @@ -0,0 +1,157 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the CHB-MIT seizure corpus into train/val/test LMDBs.""" + +import os +import pickle + +import numpy as np +from scipy import signal + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + bipolar_coordinates, + build_arg_parser, + run_jobs, +) + +SOURCE_FREQ = 256 +WINDOW_SECONDS = 10 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +SOURCE_WINDOW_SAMPLES = SOURCE_FREQ * WINDOW_SECONDS +SEIZURE_HOP_SECONDS = 5 +SEIZURE_PAD_SECONDS = 1 + +CHANNEL_NAMES = [ + "FP1-F7", "F7-T7", "T7-P7", "P7-O1", + "FP2-F8", "F8-T8", "T8-P8", "P8-O2", + "FP1-F3", "F3-C3", "C3-P3", "P3-O1", + "FP2-F4", "F4-C4", "C4-P4", "P4-O2", +] +NUM_CHANNELS = len(CHANNEL_NAMES) +CHANNEL_COORDS = bipolar_coordinates([tuple(name.split("-")) for name in CHANNEL_NAMES]) + +VAL_PATIENTS = {"chb21", "chb22"} +TEST_PATIENTS = {"chb23", "chb24"} +PATIENT_ALIASES = {"chb21": "chb01"} + + +def patient_of(path: str) -> str: + """Return the patient identifier for a recording, resolving known aliases.""" + patient = os.path.basename(os.path.dirname(path)) + return PATIENT_ALIASES.get(patient, patient) + + +def split_of(path: str) -> str: + """Assign a recording to a split by patient, holding two patients out for each.""" + patient = os.path.basename(os.path.dirname(path)) + if patient in TEST_PATIENTS: + return "test" + if patient in VAL_PATIENTS: + return "val" + return "train" + + +def process_recording(task): + """Window one recording, then re-window each seizure with overlap to enrich positives. + + Seizures are rare, so in addition to the non-overlapping pass every annotated + seizure is re-sampled with a shorter hop, extending one second either side. A + window is positive when a seizure boundary falls strictly inside it. Windows in + the seizure pass are truncated at the end of the recording and resampled to the + full window length regardless of how much signal they contain. + """ + split, path = task + with open(path, "rb") as handle: + recording = pickle.load(handle) + + try: + raw = np.stack([recording[name] for name in CHANNEL_NAMES], axis=0) + except KeyError: + return [] + + seizures = recording.get("metadata", {}).get("times", []) + record_id = os.path.basename(path).split(".")[0] + patient = patient_of(path) + length = raw.shape[1] + samples = [] + + def contains_boundary(window_start, window_end): + return any( + window_start < start < window_end or window_start < end < window_end + for start, end in seizures + ) + + def emit(key, segment, label): + resampled = signal.resample(segment, WINDOW_SAMPLES, axis=1) + samples.append(( + split, + key.encode(), + { + "eeg": resampled.astype(np.float32), + "label": int(label), + "channel_coords": CHANNEL_COORDS, + "subject_id": patient, + }, + )) + + for start in range(0, length, SOURCE_WINDOW_SAMPLES): + end = start + SOURCE_WINDOW_SAMPLES + if end > length: + continue + emit(f"{record_id}-{start}", raw[:, start:end], contains_boundary(start, end)) + + hop = SEIZURE_HOP_SECONDS * SOURCE_FREQ + pad = SEIZURE_PAD_SECONDS * SOURCE_FREQ + for seizure_index, (seizure_start, seizure_end) in enumerate(seizures): + first = max(0, seizure_start - pad) + last = min(seizure_end + pad, length) + for start in range(first, last, hop): + end = min(start + SOURCE_WINDOW_SAMPLES, length) + emit(f"{record_id}-s-{seizure_index}-add-{start}", raw[:, start:end], 1) + + return samples + + +def main(): + """Split CHB-MIT by patient and write one LMDB per split.""" + args = build_arg_parser("CHB-MIT seizure detection to LMDB").parse_args() + + recordings = [] + for directory, _, filenames in os.walk(args.input_dir): + recordings += [ + os.path.join(directory, name) for name in filenames if name.lower().endswith(".pkl") + ] + tasks = [(split_of(path), path) for path in sorted(recordings)] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "CHB-MIT recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("CHB-MIT") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_gwd.py b/make_datasets/make_gwd.py new file mode 100644 index 0000000..7b17892 --- /dev/null +++ b/make_datasets/make_gwd.py @@ -0,0 +1,115 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the GWD pre-training corpus into a pooled LMDB.""" + +import os + +import mne +import numpy as np +import scipy.io + +from make_datasets.common import ( + SAMPLING_FREQ, + build_arg_parser, + electrode_coordinate, + list_files, + slice_windows, + write_pretraining_corpus, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS +REFERENCE = "AR" + + +def known_channels(channel_names): + """Return the indices and coordinates of channels with a known electrode position.""" + reference_position = electrode_coordinate(REFERENCE) + indices, coords = [], [] + for index, name in enumerate(channel_names): + position = electrode_coordinate(name.strip().upper()) + if position == (0.0, 0.0, 0.0): + continue + indices.append(index) + coords.append([position, reference_position]) + return indices, np.asarray(coords, dtype=np.float32) + + +def process_recording(task): + """Read one MATLAB recording, keep the locatable EEG channels, filter and slice.""" + root, relative_path = task + contents = scipy.io.loadmat(os.path.join(root, relative_path)) + if "signal" not in contents or "header" not in contents: + return [] + + signal = np.asarray(contents["signal"], dtype=np.float32) + header = contents["header"] + + try: + source_freq = float(header["sample_rate"][0][0]) + except Exception: + source_freq = float(SAMPLING_FREQ) + + try: + eeg_indices = header["channels_eeg"][0][0].flatten() - 1 + raw_labels = header["channels_labels"][0][0].flatten() + labels = [str(label[0]) if isinstance(label, np.ndarray) else str(label) for label in raw_labels] + channel_names = [labels[index] for index in eeg_indices] + except Exception: + eeg_indices = np.arange(signal.shape[0]) + channel_names = [f"Ch{index + 1}" for index in eeg_indices] + + data = signal[eeg_indices, :] + info = mne.create_info(ch_names=channel_names, sfreq=source_freq, ch_types="eeg", verbose=False) + raw = mne.io.RawArray(data, info, verbose=False) + raw.notch_filter(50, verbose=False) + raw.filter(l_freq=0.5, h_freq=min(100.0, source_freq / 2 - 1), method="fir", picks="eeg", verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, verbose=False) + + indices, coords = known_channels(raw.ch_names) + if not indices: + return [] + + stem = os.path.splitext(os.path.basename(relative_path))[0] + return [ + ( + f"{stem}-{index}".encode(), + {"eeg": window, "channel_coords": coords, "subject_id": stem}, + ) + for index, window in enumerate(slice_windows(raw.get_data()[indices], WINDOW_SAMPLES)) + ] + + +def main(): + """Slice GWD into 30-second pre-training windows.""" + args = build_arg_parser("GWD pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + tasks = [(args.input_dir, name) for name in list_files(args.input_dir, [".mat"])] + write_pretraining_corpus( + "GWD", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_isruc.py b/make_datasets/make_isruc.py new file mode 100644 index 0000000..d00be2e --- /dev/null +++ b/make_datasets/make_isruc.py @@ -0,0 +1,179 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the ISRUC sleep-staging corpus into train/val/test LMDBs. + +Unlike the other fine-tuning datasets, one ISRUC sample is a sequence of consecutive +30-second epochs rather than a single window, because sleep stage depends on +neighbouring context. Each sample therefore holds ``(sequence_length, channels, +timesteps)`` with one label per epoch, and the sequence classification head consumes it. +""" + +import os + +import numpy as np +from mne.io import read_raw_edf + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + bipolar_coordinates, + build_arg_parser, + run_jobs, +) + +EPOCH_SECONDS = 30 +EPOCH_SAMPLES = SAMPLING_FREQ * EPOCH_SECONDS +SEQUENCE_LENGTH = 20 +NUM_SUBJECTS = 100 +TRAIN_MAX_SUBJECT = 80 +VAL_MAX_SUBJECT = 90 + +MASTOID_CHANNELS = ["F3-M2", "C3-M2", "O1-M2", "F4-M1", "C4-M1", "O2-M1"] +AURICULAR_CHANNELS = ["F3-A2", "C3-A2", "O1-A2", "F4-A1", "C4-A1", "O2-A1"] +NUM_CHANNELS = len(MASTOID_CHANNELS) +STAGE_TO_LABEL = {"0": 0, "1": 1, "2": 2, "3": 3, "5": 4} + +MASTOID_COORDS = bipolar_coordinates([tuple(name.split("-")) for name in MASTOID_CHANNELS]) +AURICULAR_COORDS = bipolar_coordinates([tuple(name.split("-")) for name in AURICULAR_CHANNELS]) + + +def split_of(subject: int) -> str: + """Assign a subject to a split, holding out the highest-numbered subjects for test.""" + if subject <= TRAIN_MAX_SUBJECT: + return "train" + if subject <= VAL_MAX_SUBJECT: + return "val" + return "test" + + +def recording_pairs(input_dir: str): + """Return the recording and hypnogram paths for every available subject.""" + pairs = [] + for subject in range(1, NUM_SUBJECTS + 1): + folder = os.path.join(input_dir, str(subject)) + recording = os.path.join(folder, f"{subject}.rec") + hypnogram = os.path.join(folder, f"{subject}_1.txt") + if os.path.isfile(recording) and os.path.isfile(hypnogram): + pairs.append((subject, recording, hypnogram)) + return pairs + + +def resolve_channels(recording: str): + """Choose the montage present in a recording, deriving it from unipolar channels if needed. + + Returns the channel names to pick and their coordinates, or ``(None, None, None)`` when + the recording has none of the expected montages. + """ + present = set(read_raw_edf(recording, preload=False, verbose=False).info["ch_names"]) + if all(name in present for name in MASTOID_CHANNELS): + return MASTOID_CHANNELS, MASTOID_COORDS, False + if all(name in present for name in AURICULAR_CHANNELS): + return AURICULAR_CHANNELS, AURICULAR_COORDS, False + electrodes = {name for channel in AURICULAR_CHANNELS for name in channel.split("-")} + if electrodes <= present: + return sorted(electrodes), AURICULAR_COORDS, True + return None, None, None + + +def process_recording(task): + """Cut one night into sequences of consecutive labelled epochs.""" + split, subject, recording, hypnogram, channels, coords, derive_bipolar = task + + raw = read_raw_edf(recording, preload=True, verbose=False) + raw.filter(0.3, 35, fir_design="firwin", verbose=False) + raw.notch_filter(50, verbose=False) + raw.pick(channels) + raw.reorder_channels(channels) + data = raw.to_data_frame().values[:, 1:].T + + if derive_bipolar: + pairs = [name.split("-") for name in AURICULAR_CHANNELS] + derived = np.zeros((NUM_CHANNELS, data.shape[1]), dtype=np.float32) + for index, (active, reference) in enumerate(pairs): + derived[index] = data[channels.index(active)] - data[channels.index(reference)] + data = derived + + transposed = data.T + remainder = transposed.shape[0] % EPOCH_SAMPLES + if remainder: + transposed = transposed[:-remainder] + epochs = transposed.reshape(-1, EPOCH_SAMPLES, NUM_CHANNELS) + + with open(hypnogram) as handle: + labels = np.array([STAGE_TO_LABEL[line.strip()] for line in handle if line.strip()]) + + usable = min(epochs.shape[0], labels.shape[0]) + epochs, labels = epochs[:usable], labels[:usable] + + trailing = epochs.shape[0] % SEQUENCE_LENGTH + if trailing: + epochs, labels = epochs[:-trailing], labels[:-trailing] + if epochs.shape[0] == 0: + return [] + + num_sequences = epochs.shape[0] // SEQUENCE_LENGTH + epochs = epochs.reshape(num_sequences, SEQUENCE_LENGTH, EPOCH_SAMPLES, NUM_CHANNELS) + epochs = epochs.transpose(0, 1, 3, 2) + labels = labels.reshape(num_sequences, SEQUENCE_LENGTH) + + return [ + ( + split, + f"{subject}-{index}".encode(), + { + "eeg": epochs[index].astype(np.float32), + "label": labels[index].tolist(), + "channel_coords": coords, + "subject_id": str(subject), + }, + ) + for index in range(num_sequences) + ] + + +def main(): + """Split ISRUC by subject and write one LMDB per split.""" + args = build_arg_parser("ISRUC sleep staging to LMDB").parse_args() + + tasks, skipped = [], [] + for subject, recording, hypnogram in recording_pairs(args.input_dir): + channels, coords, derive_bipolar = resolve_channels(recording) + if channels is None: + skipped.append(subject) + continue + tasks.append( + (split_of(subject), subject, recording, hypnogram, channels, coords, derive_bipolar) + ) + + print(f"[ISRUC] usable subjects: {len(tasks)}, skipped for missing channels: {len(skipped)}") + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "ISRUC nights"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("ISRUC") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_mental_arithmetic.py b/make_datasets/make_mental_arithmetic.py new file mode 100644 index 0000000..5d91154 --- /dev/null +++ b/make_datasets/make_mental_arithmetic.py @@ -0,0 +1,145 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the mental arithmetic corpus into train/val/test LMDBs.""" + +import os +import re + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + electrode_coordinate, + index_splits, + run_jobs, +) + +WINDOW_SECONDS = 5 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +TRAIN_RECORDINGS = 56 +VAL_RECORDINGS = 64 + +EDF_CHANNELS = [ + "EEG Fp1", "EEG Fp2", "EEG F3", "EEG F4", "EEG F7", "EEG F8", + "EEG T3", "EEG T4", "EEG C3", "EEG C4", "EEG T5", "EEG T6", + "EEG P3", "EEG P4", "EEG O1", "EEG O2", "EEG Fz", "EEG Cz", + "EEG Pz", "EEG A2-A1", +] +NUM_CHANNELS = len(EDF_CHANNELS) +REFERENCE = "A1" + + +def channel_coordinates() -> np.ndarray: + """Build coordinates for the montage, treating the A2-A1 channel as bipolar.""" + coords = np.zeros((NUM_CHANNELS, 2, 3), dtype=np.float32) + reference_position = electrode_coordinate(REFERENCE) + for index, name in enumerate(EDF_CHANNELS): + label = name.replace("EEG ", "").upper() + if label == "A2-A1": + coords[index, 0, :] = electrode_coordinate("A2") + coords[index, 1, :] = electrode_coordinate("A1") + else: + coords[index, 0, :] = electrode_coordinate(label) + coords[index, 1, :] = reference_position + return coords + + +CHANNEL_COORDS = channel_coordinates() + +SUBJECT_PATTERN = re.compile(r"Subject(\d+)", re.IGNORECASE) +LABEL_PATTERN = re.compile(r"_(\d)\.edf$", re.IGNORECASE) + + +def subject_of(filename: str) -> str: + """Return the subject number encoded in a recording filename.""" + match = SUBJECT_PATTERN.search(filename) + return (match.group(1).lstrip("0") or "0") if match else os.path.splitext(filename)[0] + + +def label_of(filename: str) -> int: + """Return the task condition encoded in the filename suffix.""" + match = LABEL_PATTERN.search(os.path.basename(filename)) + if not match: + raise ValueError(f"Cannot parse a condition label from {filename}") + return int(match.group(1)) - 1 + + +def process_recording(task): + """Resample one recording and cut it into fixed-length windows.""" + split, root, filename = task + raw = mne.io.read_raw_edf(os.path.join(root, filename), preload=True, verbose=False) + raw.pick(EDF_CHANNELS) + raw.reorder_channels(EDF_CHANNELS) + raw.resample(SAMPLING_FREQ, verbose=False) + + data = raw.get_data(units="uV") + if data.shape[0] != NUM_CHANNELS: + raise RuntimeError(f"{filename} has {data.shape[0]} channels, expected {NUM_CHANNELS}") + + num_windows = data.shape[1] // WINDOW_SAMPLES + if num_windows == 0: + return [] + windows = data[:, : num_windows * WINDOW_SAMPLES] + windows = windows.reshape(NUM_CHANNELS, num_windows, WINDOW_SAMPLES).transpose(1, 0, 2) + + subject = subject_of(filename) + label = label_of(filename) + stem = os.path.splitext(filename)[0] + return [ + ( + split, + f"{stem}-{index}".encode(), + { + "eeg": window.astype(np.float32), + "label": label, + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + ) + for index, window in enumerate(windows) + ] + + +def main(): + """Split the mental arithmetic corpus by recording and write one LMDB per split.""" + args = build_arg_parser("Mental arithmetic classification to LMDB").parse_args() + + recordings = sorted(name for name in os.listdir(args.input_dir) if name.lower().endswith(".edf")) + splits = index_splits(recordings, TRAIN_RECORDINGS, VAL_RECORDINGS) + tasks = [ + (split, args.input_dir, filename) for split, names in splits.items() for filename in names + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "Mental arithmetic recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("MentalArithmetic") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_mumtaz.py b/make_datasets/make_mumtaz.py new file mode 100644 index 0000000..9c2ada9 --- /dev/null +++ b/make_datasets/make_mumtaz.py @@ -0,0 +1,147 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the Mumtaz depression corpus into train/val/test LMDBs.""" + +import os +import re + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + referential_coordinates, + run_jobs, +) + +WINDOW_SECONDS = 5 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +REFERENCE = "LE" + +EDF_CHANNELS = [ + "EEG Fp1-LE", "EEG Fp2-LE", "EEG F3-LE", "EEG F4-LE", "EEG C3-LE", "EEG C4-LE", + "EEG P3-LE", "EEG P4-LE", "EEG O1-LE", "EEG O2-LE", "EEG F7-LE", "EEG F8-LE", + "EEG T3-LE", "EEG T4-LE", "EEG T5-LE", "EEG T6-LE", "EEG Fz-LE", "EEG Cz-LE", + "EEG Pz-LE", +] +CHANNELS = [name.replace("EEG ", "").replace("-LE", "").upper() for name in EDF_CHANNELS] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, REFERENCE) + +HEALTHY_TRAIN, HEALTHY_VAL = 40, 48 +DEPRESSED_TRAIN, DEPRESSED_VAL = 42, 52 + +SUBJECT_PATTERN = re.compile(r"\b(H|MDD)\s*S?(\d+)", re.IGNORECASE) + + +def subject_of(filename: str): + """Return the group-qualified subject identifier, or None if it cannot be parsed.""" + match = SUBJECT_PATTERN.search(os.path.splitext(os.path.basename(filename))[0].replace(" ", " ")) + if not match: + return None + return f"{match.group(1).upper()}_S{match.group(2)}" + + +def label_of(subject: str) -> int: + """Return 1 for depressed subjects and 0 for healthy controls.""" + return 1 if subject.startswith("MDD_") else 0 + + +def process_recording(task): + """Filter one resting-state recording and cut it into fixed-length windows.""" + split, root, filename = task + if "TASK" in filename.upper(): + return [] + subject = subject_of(filename) + if subject is None: + return [] + + raw = mne.io.read_raw_edf(os.path.join(root, filename), preload=True, verbose=False) + present = [name for name in EDF_CHANNELS if name in raw.ch_names] + if len(present) != NUM_CHANNELS: + return [] + raw.pick(present) + raw.reorder_channels(present) + raw.resample(SAMPLING_FREQ, verbose=False) + raw.filter(l_freq=0.3, h_freq=75, verbose=False) + raw.notch_filter(50, verbose=False) + + data = raw.get_data(units="uV") + num_windows = data.shape[1] // WINDOW_SAMPLES + if num_windows == 0: + return [] + + windows = data[:, : num_windows * WINDOW_SAMPLES] + windows = windows.reshape(NUM_CHANNELS, num_windows, WINDOW_SAMPLES).transpose(1, 0, 2) + + label = label_of(subject) + stem = os.path.splitext(filename)[0] + return [ + ( + split, + f"{stem}-{index}".encode(), + { + "eeg": window.astype(np.float32), + "label": label, + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + ) + for index, window in enumerate(windows) + ] + + +def main(): + """Split Mumtaz by recording within each diagnostic group and write one LMDB per split.""" + args = build_arg_parser("Mumtaz depression detection to LMDB").parse_args() + + healthy, depressed = [], [] + for filename in sorted(os.listdir(args.input_dir)): + if not filename.lower().endswith(".edf") or "TASK" in filename.upper(): + continue + subject = subject_of(filename) + if subject is None: + continue + (depressed if label_of(subject) else healthy).append(filename) + + splits = { + "train": healthy[:HEALTHY_TRAIN] + depressed[:DEPRESSED_TRAIN], + "val": healthy[HEALTHY_TRAIN:HEALTHY_VAL] + depressed[DEPRESSED_TRAIN:DEPRESSED_VAL], + "test": healthy[HEALTHY_VAL:] + depressed[DEPRESSED_VAL:], + } + tasks = [ + (split, args.input_dir, filename) for split, names in splits.items() for filename in names + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "Mumtaz recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("Mumtaz") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_neonate.py b/make_datasets/make_neonate.py new file mode 100644 index 0000000..54c219b --- /dev/null +++ b/make_datasets/make_neonate.py @@ -0,0 +1,207 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the Helsinki neonatal seizure corpus into train/val/test LMDBs.""" + +import os +import re + +import mne +import numpy as np +import pandas as pd + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + electrode_coordinate, + run_jobs, +) + +WINDOW_SECONDS = 5 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +DROPPED_CHANNELS = ["ECG EKG", "Resp Effort", "ECG EKG-REF", "Resp Effort-REF"] +ANNOTATION_FILES = ["annotations_2017_A.csv", "annotations_2017_B.csv", "annotations_2017_C.csv"] +TIE_BREAK_SEED = 0 + +BIPOLAR_MONTAGE = [ + ("Fp2", "F4"), ("F4", "C4"), ("C4", "P4"), ("P4", "O2"), + ("Fp1", "F3"), ("F3", "C3"), ("C3", "P3"), ("P3", "O1"), + ("Fp2", "F8"), ("F8", "T4"), ("T4", "T6"), ("T6", "O2"), + ("Fp1", "F7"), ("F7", "T3"), ("T3", "T5"), ("T5", "O1"), + ("Fz", "Cz"), ("Cz", "Pz"), +] +NUM_CHANNELS = len(BIPOLAR_MONTAGE) + +SUBJECT_SPLITS = { + "train": {2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 25, 26, 28, + 29, 30, 31, 33, 35, 36, 37, 38, 40, 41, 43, 44, 45, 48, 49, 50, 51, 52, 54, 55, + 56, 58, 60, 61, 62, 63, 64, 65, 66, 67, 70, 71, 72, 73, 74, 75, 76, 77, 78}, + "val": {1, 3, 24, 32, 34, 42, 46, 69}, + "test": {7, 11, 27, 39, 47, 53, 57, 59, 68, 79}, +} + +SUBJECT_PATTERN = re.compile(r"eeg(\d+)\.edf$", re.IGNORECASE) + + +def subject_of(filename: str): + """Return the subject number encoded in a recording filename, or None.""" + match = SUBJECT_PATTERN.search(filename) + return int(match.group(1)) if match else None + + +def split_of(subject: int): + """Return the split a subject belongs to, or None if it is not used.""" + for split, subjects in SUBJECT_SPLITS.items(): + if subject in subjects: + return split + return None + + +def consensus_annotations(input_dir: str) -> pd.DataFrame: + """Combine the three expert annotation files into a per-second consensus label. + + A second is labelled when at least two annotators agree. Two-annotator ties are + broken with a seeded draw, and seconds without a majority are left as NaN so the + window is skipped. + """ + tables = [pd.read_csv(os.path.join(input_dir, name), header=None) for name in ANNOTATION_FILES] + if len({table.shape for table in tables}) != 1: + raise ValueError("Annotation files must have identical shapes") + + generator = np.random.default_rng(TIE_BREAK_SEED) + + def consensus(first, second, third): + values = np.array([first, second, third], dtype=float) + present = values[~np.isnan(values)] + if present.size == 0: + return np.nan + if present.size == 1: + return present[0] + if present.size == 2: + return present[0] if present[0] == present[1] else float(generator.integers(2)) + return float(round(present.mean())) + + combined = np.vectorize(consensus)(*[table.values for table in tables]) + frame = pd.DataFrame(combined) + return frame.where(frame.isin([0.0, 1.0]), np.nan) + + +def bipolar_montage(raw): + """Derive the bipolar montage from a referential recording. + + Returns the montage signal and its channel coordinates, keeping only the pairs + whose electrodes are both present in the recording. + """ + available = {} + for name in raw.ch_names: + upper = name.upper() + if upper.startswith("EEG ") and "-REF" in upper: + available[upper[4:].split("-")[0].strip()] = name + + data = raw.get_data() + channel_names = list(raw.ch_names) + signals, coords = [], [] + + for active, reference in BIPOLAR_MONTAGE: + active_channel = available.get(active.upper()) + reference_channel = available.get(reference.upper()) + if active_channel is None or reference_channel is None: + continue + signals.append(data[channel_names.index(active_channel)] - data[channel_names.index(reference_channel)]) + coords.append([electrode_coordinate(active.upper()), electrode_coordinate(reference.upper())]) + + if not signals: + return None, None + return np.stack(signals), np.asarray(coords, dtype=np.float32) + + +def process_recording(task): + """Filter one recording, build its bipolar montage and emit one window per second.""" + split, input_dir, filename, annotations = task + subject = subject_of(filename) + if subject is None: + return [] + + raw = mne.io.read_raw_edf(os.path.join(input_dir, filename), preload=True, verbose="ERROR") + kept = [name for name in raw.ch_names if name not in DROPPED_CHANNELS] + if not kept: + return [] + raw.pick(kept) + raw.filter(l_freq=0.5, h_freq=None, method="iir", phase="forward", + iir_params=dict(order=6, ftype="butter"), verbose="ERROR") + raw.notch_filter(freqs=50, notch_widths=4.0, method="iir", phase="forward", verbose="ERROR") + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, n_jobs=1, verbose="ERROR") + + montage, coords = bipolar_montage(raw) + if montage is None or montage.shape[1] < WINDOW_SAMPLES: + return [] + + labels = annotations.iloc[:, subject - 1].values + seconds = min(int(montage.shape[1] // SAMPLING_FREQ), len(labels)) + stem = os.path.splitext(filename)[0] + + samples = [] + for second in range(seconds): + label = labels[second] + if label not in (0, 1): + continue + start = second * WINDOW_SAMPLES + window = montage[:, start : start + WINDOW_SAMPLES] + if window.shape[1] != WINDOW_SAMPLES: + continue + samples.append(( + split, + f"{stem}-{second}".encode(), + { + "eeg": window.astype(np.float32), + "label": int(label), + "channel_coords": coords, + "subject_id": str(subject), + }, + )) + return samples + + +def main(): + """Split the neonatal corpus by subject and write one LMDB per split.""" + args = build_arg_parser("Neonatal seizure detection to LMDB").parse_args() + annotations = consensus_annotations(args.input_dir) + + tasks = [] + for filename in sorted(name for name in os.listdir(args.input_dir) if name.lower().endswith(".edf")): + subject = subject_of(filename) + split = split_of(subject) if subject is not None else None + if split is not None: + tasks.append((split, args.input_dir, filename, annotations)) + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "Neonate recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("Neonate") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_physionet_mi.py b/make_datasets/make_physionet_mi.py new file mode 100644 index 0000000..88a75c6 --- /dev/null +++ b/make_datasets/make_physionet_mi.py @@ -0,0 +1,135 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the PhysioNet motor imagery corpus into train/val/test LMDBs.""" + +import os + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + index_splits, + referential_coordinates, + run_jobs, +) + +WINDOW_SECONDS = 4 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +RUNS = ["04", "06", "08", "10", "12", "14"] +TRAIN_SUBJECTS = 70 +VAL_SUBJECTS = 89 +TEST_SUBJECTS = 109 + +EDF_CHANNELS = [ + "Fc5.", "Fc3.", "Fc1.", "Fcz.", "Fc2.", "Fc4.", "Fc6.", + "C5..", "C3..", "C1..", "Cz..", "C2..", "C4..", "C6..", + "Cp5.", "Cp3.", "Cp1.", "Cpz.", "Cp2.", "Cp4.", "Cp6.", + "Fp1.", "Fpz.", "Fp2.", "Af7.", "Af3.", "Afz.", "Af4.", "Af8.", + "F7..", "F5..", "F3..", "F1..", "Fz..", "F2..", "F4..", "F6..", "F8..", + "Ft7.", "Ft8.", "T7..", "T8..", "T9..", "T10.", + "Tp7.", "Tp8.", + "P7..", "P5..", "P3..", "P1..", "Pz..", "P2..", "P4..", "P6..", "P8..", + "Po7.", "Po3.", "Poz.", "Po4.", "Po8.", + "O1..", "Oz..", "O2..", "Iz..", +] +CHANNELS = [name.strip(".").upper() for name in EDF_CHANNELS] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + + +def event_to_label(event: int, run: str) -> int: + """Map an annotation code to a class index, accounting for the run's task pairing.""" + return int(event - 2) if run in ("04", "08", "12") else int(event) + + +def process_run(task): + """Epoch one motor imagery run and drop rest epochs.""" + split, root, subject, run = task + path = os.path.join(root, subject, f"{subject}R{run}.edf") + if not os.path.isfile(path): + return [] + + raw = mne.io.read_raw_edf(path, preload=True, verbose=False) + raw.pick(EDF_CHANNELS) + raw.reorder_channels(EDF_CHANNELS) + if raw.info["bads"]: + raw.interpolate_bads() + raw.set_eeg_reference(ref_channels="average", verbose=False) + raw.filter(l_freq=0.3, h_freq=None, verbose=False) + raw.notch_filter(60, verbose=False) + raw.resample(SAMPLING_FREQ, verbose=False) + + events, event_ids = mne.events_from_annotations(raw, verbose=False) + epochs = mne.Epochs( + raw, events, event_ids, tmin=0, tmax=WINDOW_SECONDS - 1.0 / SAMPLING_FREQ, + baseline=None, preload=True, verbose=False, + ) + data = epochs.get_data(units="uV")[:, :, -WINDOW_SAMPLES:] + + samples = [] + for index, (window, event) in enumerate(zip(data, epochs.events[:, 2])): + if event == 1: + continue + samples.append(( + split, + f"{subject}R{run}-{index}".encode(), + { + "eeg": window.astype(np.float32), + "label": event_to_label(event, run), + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + )) + return samples + + +def main(): + """Split PhysioNet-MI by subject and write one LMDB per split.""" + args = build_arg_parser("PhysioNet motor imagery to LMDB").parse_args() + + subjects = sorted( + name for name in os.listdir(args.input_dir) + if os.path.isdir(os.path.join(args.input_dir, name)) + )[:TEST_SUBJECTS] + splits = index_splits(subjects, TRAIN_SUBJECTS, VAL_SUBJECTS) + + tasks = [ + (split, args.input_dir, subject, run) + for split, names in splits.items() + for subject in names + for run in RUNS + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_run, tasks, args.num_workers, "PhysioNet-MI runs"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("PhysioNet-MI") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_seed.py b/make_datasets/make_seed.py new file mode 100644 index 0000000..809740d --- /dev/null +++ b/make_datasets/make_seed.py @@ -0,0 +1,91 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the SEED and SEED-IV pre-training corpora into a pooled LMDB.""" + +import os + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + build_arg_parser, + referential_coordinates, + slice_windows, + write_pretraining_corpus, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS + +CHANNELS = [ + "FP1", "FPZ", "FP2", "AF3", "AF4", "F7", "F5", "F3", "F1", "FZ", "F2", "F4", "F6", "F8", + "FT7", "FC5", "FC3", "FC1", "FCZ", "FC2", "FC4", "FC6", "FT8", "T7", "C5", "C3", "C1", + "CZ", "C2", "C4", "C6", "T8", "TP7", "CP5", "CP3", "CP1", "CPZ", "CP2", "CP4", "CP6", + "TP8", "P7", "P5", "P3", "P1", "PZ", "P2", "P4", "P6", "P8", "PO7", "PO5", "PO3", "POZ", + "PO4", "PO6", "PO8", "CB1", "O1", "OZ", "O2", "CB2", +] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + + +def process_recording(task): + """Slice every 62-channel array stored in one MATLAB session file.""" + root, relative_path = task + import scipy.io + + contents = scipy.io.loadmat(os.path.join(root, relative_path)) + stem = os.path.splitext(os.path.basename(relative_path))[0] + samples = [] + + for name, value in contents.items(): + if name.startswith("__"): + continue + array = np.asarray(value, dtype=np.float32) + if array.ndim != 2 or array.shape[0] != NUM_CHANNELS: + continue + + info = mne.create_info(CHANNELS, sfreq=SAMPLING_FREQ, ch_types="eeg", verbose=False) + raw = mne.io.RawArray(array, info, verbose=False) + for index, window in enumerate(slice_windows(raw.get_data(), WINDOW_SAMPLES)): + samples.append(( + f"{stem}-{name}-{index}".encode(), + {"eeg": window, "channel_coords": CHANNEL_COORDS, "subject_id": stem}, + )) + return samples + + +def main(): + """Slice SEED or SEED-IV into 30-second pre-training windows.""" + args = build_arg_parser("SEED / SEED-IV pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + from make_datasets.common import list_files + + tasks = [(args.input_dir, name) for name in list_files(args.input_dir, [".mat"])] + write_pretraining_corpus( + "SEED", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_seed_fra_ger.py b/make_datasets/make_seed_fra_ger.py new file mode 100644 index 0000000..a0cb4cf --- /dev/null +++ b/make_datasets/make_seed_fra_ger.py @@ -0,0 +1,93 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the SEED-FRA and SEED-GER pre-training corpora into a pooled LMDB.""" + +import os + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + build_arg_parser, + referential_coordinates, + slice_windows, + write_pretraining_corpus, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS + +CHANNELS = [ + "FP1", "FPZ", "FP2", "AF3", "AF4", "F7", "F5", "F3", "F1", "FZ", "F2", "F4", "F6", "F8", + "FT7", "FC5", "FC3", "FC1", "FCZ", "FC2", "FC4", "FC6", "FT8", "T7", "C5", "C3", "C1", + "CZ", "C2", "C4", "C6", "T8", "TP7", "CP5", "CP3", "CP1", "CPZ", "CP2", "CP4", "CP6", + "TP8", "P7", "P5", "P3", "P1", "PZ", "P2", "P4", "P6", "P8", "PO7", "PO5", "PO3", "POZ", + "PO4", "PO6", "PO8", "CB1", "O1", "OZ", "O2", "CB2", +] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + + +def process_recording(task): + """Filter and slice one Neuroscan session recording.""" + root, relative_path = task + path = os.path.join(root, relative_path) + raw = mne.io.read_raw_cnt(path, preload=True, verbose=False) + if raw.n_times == 0: + return [] + + present = [name for name in CHANNELS if name in raw.ch_names] + if len(present) != NUM_CHANNELS: + return [] + raw.pick(present) + raw.reorder_channels(CHANNELS) + raw.notch_filter(50, verbose=False) + raw.filter(l_freq=0.5, h_freq=30, method="fir", picks="eeg", verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, verbose=False) + + stem = os.path.splitext(os.path.basename(relative_path))[0] + return [ + ( + f"{stem}-{index}".encode(), + {"eeg": window, "channel_coords": CHANNEL_COORDS, "subject_id": stem}, + ) + for index, window in enumerate(slice_windows(raw.get_data(), WINDOW_SAMPLES)) + ] + + +def main(): + """Slice SEED-FRA or SEED-GER into 30-second pre-training windows.""" + args = build_arg_parser("SEED-FRA / SEED-GER pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + from make_datasets.common import list_files + + tasks = [(args.input_dir, name) for name in list_files(args.input_dir, [".cnt"])] + write_pretraining_corpus( + "SEED-FRA", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_seed_v.py b/make_datasets/make_seed_v.py new file mode 100644 index 0000000..f51f6cc --- /dev/null +++ b/make_datasets/make_seed_v.py @@ -0,0 +1,149 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the SEED-V emotion corpus into train/val/test LMDBs.""" + +import os +import re + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + referential_coordinates, + run_jobs, +) + +SEGMENT_SECONDS = 4 +SEGMENT_SAMPLES = SAMPLING_FREQ * SEGMENT_SECONDS +TRIALS_PER_SESSION = 15 + +CHANNELS = [ + "FP1", "FPZ", "FP2", "AF3", "AF4", "F7", "F5", "F3", "F1", "FZ", "F2", "F4", "F6", "F8", + "FT7", "FC5", "FC3", "FC1", "FCZ", "FC2", "FC4", "FC6", "FT8", "T7", "C5", "C3", "C1", + "CZ", "C2", "C4", "C6", "T8", "TP7", "CP5", "CP3", "CP1", "CPZ", "CP2", "CP4", "CP6", + "TP8", "P7", "P5", "P3", "P1", "PZ", "P2", "P4", "P6", "P8", "PO7", "PO5", "PO3", "POZ", + "PO4", "PO6", "PO8", "CB1", "O1", "OZ", "O2", "CB2", +] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + +TRIAL_BOUNDS = { + "1": {"start": [30, 132, 287, 555, 773, 982, 1271, 1628, 1730, 2025, 2227, 2435, 2667, 2932, 3204], + "end": [102, 228, 524, 742, 920, 1240, 1568, 1697, 1994, 2166, 2401, 2607, 2901, 3172, 3359]}, + "2": {"start": [30, 299, 548, 646, 836, 1000, 1091, 1392, 1657, 1809, 1966, 2186, 2333, 2490, 2741], + "end": [267, 488, 614, 773, 967, 1059, 1331, 1622, 1777, 1908, 2153, 2302, 2428, 2709, 2817]}, + "3": {"start": [30, 353, 478, 674, 825, 908, 1200, 1346, 1451, 1711, 2055, 2307, 2457, 2726, 2888], + "end": [321, 418, 643, 764, 877, 1147, 1284, 1418, 1679, 1996, 2275, 2425, 2664, 2857, 3066]}, +} + +TRIAL_LABELS = { + "1": [4, 1, 3, 2, 0, 4, 1, 3, 2, 0, 4, 1, 3, 2, 0], + "2": [2, 1, 3, 0, 4, 4, 0, 3, 2, 1, 3, 4, 1, 2, 0], + "3": [2, 1, 3, 0, 4, 4, 0, 3, 2, 1, 3, 4, 1, 2, 0], +} + +SUBJECT_SPLITS = { + "train": {str(index) for index in range(1, 11)}, + "val": {str(index) for index in range(11, 14)}, + "test": {str(index) for index in range(14, 17)}, +} + +SUBJECT_PATTERN = re.compile(r"(\d+)") + + +def subject_of(filename: str) -> str: + """Return the subject number encoded at the start of a SEED-V filename.""" + match = SUBJECT_PATTERN.search(os.path.basename(filename)) + return match.group(1) if match else os.path.splitext(os.path.basename(filename))[0] + + +def split_of(subject: str): + """Return the split a subject belongs to, or None if it is not used.""" + for split, subjects in SUBJECT_SPLITS.items(): + if subject in subjects: + return split + return None + + +def process_recording(task): + """Cut each labelled trial of one session into fixed-length segments.""" + root, filename = task + subject = subject_of(filename) + split = split_of(subject) + if split is None: + return [] + + raw = mne.io.read_raw_cnt(os.path.join(root, filename), preload=True, verbose=False) + raw.pick(CHANNELS) + raw.reorder_channels(CHANNELS) + raw.resample(SAMPLING_FREQ, verbose=False) + raw.filter(l_freq=0.3, h_freq=75, verbose=False) + data = raw.get_data(units="uV") + + session = os.path.basename(filename).split("_")[1] + bounds = TRIAL_BOUNDS[session] + labels = TRIAL_LABELS[session] + + samples = [] + for trial in range(TRIALS_PER_SESSION): + trial_data = data[:, bounds["start"][trial] * SAMPLING_FREQ : bounds["end"][trial] * SAMPLING_FREQ] + num_segments = trial_data.shape[1] // SEGMENT_SAMPLES + if num_segments == 0: + continue + trimmed = trial_data[:, : num_segments * SEGMENT_SAMPLES] + segments = trimmed.reshape(NUM_CHANNELS, num_segments, SEGMENT_SAMPLES).transpose(1, 0, 2) + for index, segment in enumerate(segments): + samples.append(( + split, + f"{os.path.basename(filename)[:-4]}-{trial}-{index}".encode(), + { + "eeg": segment.astype(np.float32), + "label": int(labels[trial]), + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + )) + return samples + + +def main(): + """Split SEED-V by subject and write one LMDB per split.""" + args = build_arg_parser("SEED-V emotion recognition to LMDB").parse_args() + mne.set_log_level("ERROR") + + recordings = sorted(name for name in os.listdir(args.input_dir) if name.lower().endswith(".cnt")) + tasks = [(args.input_dir, name) for name in recordings] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "SEED-V sessions"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("SEED-V") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_seed_vig.py b/make_datasets/make_seed_vig.py new file mode 100644 index 0000000..0de5b64 --- /dev/null +++ b/make_datasets/make_seed_vig.py @@ -0,0 +1,119 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the SEED-VIG vigilance corpus into train/val/test LMDBs.""" + +import os +import re + +import numpy as np +import scipy.io + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + index_splits, + referential_coordinates, + run_jobs, +) + +WINDOW_SECONDS = 8 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +TRAIN_RECORDINGS = 15 +VAL_RECORDINGS = 19 + +CHANNELS = [ + "FT7", "FT8", "T7", "T8", "TP7", "TP8", "CP1", "CP2", + "P1", "PZ", "P2", "PO3", "POZ", "PO4", "O1", "OZ", "O2", +] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + +SUBJECT_PATTERN = re.compile(r"(?:sub(?:ject)?[_\- ]?|S)(\d+)", re.IGNORECASE) + + +def subject_of(filename: str) -> str: + """Return the subject number encoded in a recording filename.""" + match = SUBJECT_PATTERN.search(os.path.basename(filename)) + return match.group(1) if match else os.path.splitext(os.path.basename(filename))[0] + + +def process_recording(task): + """Pair each 8-second window of one session with its PERCLOS score.""" + split, data_dir, labels_dir, filename = task + eeg = scipy.io.loadmat(os.path.join(data_dir, filename))["EEG"][0][0][0] + labels = scipy.io.loadmat(os.path.join(labels_dir, filename))["perclos"][:, 0] + + total_points, channels = eeg.shape + if channels != NUM_CHANNELS: + raise ValueError(f"{filename} has {channels} channels, expected {NUM_CHANNELS}") + + num_windows = total_points // WINDOW_SAMPLES + if total_points != num_windows * WINDOW_SAMPLES: + raise ValueError(f"{filename} length {total_points} is not a multiple of {WINDOW_SAMPLES}") + if len(labels) != num_windows: + raise ValueError(f"{filename} has {len(labels)} labels for {num_windows} windows") + + windows = eeg.reshape(num_windows, WINDOW_SAMPLES, NUM_CHANNELS).transpose(0, 2, 1) + subject = subject_of(filename) + stem = os.path.splitext(os.path.basename(filename))[0] + + return [ + ( + split, + f"{stem}-{index}".encode(), + { + "eeg": window.astype(np.float32), + "label": float(label), + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + ) + for index, (window, label) in enumerate(zip(windows, labels)) + ] + + +def main(): + """Split SEED-VIG by recording and write one LMDB per split.""" + parser = build_arg_parser("SEED-VIG vigilance regression to LMDB") + parser.add_argument("--labels_dir", required=True, help="Directory holding the PERCLOS label files") + args = parser.parse_args() + + recordings = sorted(name for name in os.listdir(args.input_dir) if name.lower().endswith(".mat")) + splits = index_splits(recordings, TRAIN_RECORDINGS, VAL_RECORDINGS) + tasks = [ + (split, args.input_dir, args.labels_dir, filename) + for split, names in splits.items() + for filename in names + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "SEED-VIG sessions"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("SEED-VIG") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_shu_mi.py b/make_datasets/make_shu_mi.py new file mode 100644 index 0000000..d4346df --- /dev/null +++ b/make_datasets/make_shu_mi.py @@ -0,0 +1,99 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the SHU-MI motor imagery corpus into train/val/test LMDBs.""" + +import os + +import numpy as np +import scipy.io +from scipy import signal + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + index_splits, + list_files, + referential_coordinates, + run_jobs, +) + +WINDOW_SECONDS = 4 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +TRAIN_RECORDINGS = 75 +VAL_RECORDINGS = 100 + +CHANNELS = [ + "FP1", "FP2", "FZ", "F3", "F4", "F7", "F8", "FC1", "FC2", "FC5", "FC6", + "CZ", "C3", "C4", "T3", "T4", "A1", "A2", "CP1", "CP2", "CP5", "CP6", + "PZ", "P3", "P4", "T5", "T6", "PO3", "PO4", "OZ", "O1", "O2", +] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + + +def process_recording(task): + """Resample one session to the target rate and emit one sample per trial.""" + split, root, relative_path = task + data = scipy.io.loadmat(os.path.join(root, relative_path)) + trials = signal.resample(data["data"], WINDOW_SAMPLES, axis=2) + labels = data["labels"][0] + stem = os.path.splitext(os.path.basename(relative_path))[0] + + samples = [] + for index in range(trials.shape[0]): + samples.append(( + split, + f"{stem}-{index}".encode(), + { + "eeg": trials[index].astype(np.float32), + "label": int(labels[index] - 1), + "channel_coords": CHANNEL_COORDS, + "subject_id": stem, + }, + )) + return samples + + +def main(): + """Split SHU-MI by recording and write one LMDB per split.""" + args = build_arg_parser("SHU-MI motor imagery to LMDB").parse_args() + + recordings = list_files(args.input_dir, [".mat"]) + splits = index_splits(recordings, TRAIN_RECORDINGS, VAL_RECORDINGS) + tasks = [ + (split, args.input_dir, relative_path) + for split, paths in splits.items() + for relative_path in paths + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "SHU-MI recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("SHU-MI") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_sleepedfx.py b/make_datasets/make_sleepedfx.py new file mode 100644 index 0000000..91ff8f9 --- /dev/null +++ b/make_datasets/make_sleepedfx.py @@ -0,0 +1,86 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the Sleep-EDFx pre-training corpus into a pooled LMDB.""" + +import os + +import mne + +from make_datasets.common import ( + SAMPLING_FREQ, + bipolar_coordinates, + build_arg_parser, + list_files, + slice_windows, + write_pretraining_corpus, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS + +EDF_CHANNELS = ["EEG Fpz-Cz", "EEG Pz-Oz"] +NUM_CHANNELS = len(EDF_CHANNELS) +CHANNEL_COORDS = bipolar_coordinates([ + tuple(name.replace("EEG ", "").upper().split("-")) for name in EDF_CHANNELS +]) + + +def process_recording(task): + """Filter and slice one polysomnography recording.""" + root, relative_path = task + raw = mne.io.read_raw_edf(os.path.join(root, relative_path), preload=True, verbose=False) + + present = [name for name in EDF_CHANNELS if name in raw.ch_names] + if len(present) != NUM_CHANNELS: + return [] + raw.pick(present) + raw.reorder_channels(EDF_CHANNELS) + raw.filter(l_freq=0.5, h_freq=30, method="fir", picks="eeg", verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, verbose=False) + + stem = os.path.splitext(os.path.basename(relative_path))[0] + return [ + ( + f"{stem}-{index}".encode(), + {"eeg": window, "channel_coords": CHANNEL_COORDS, "subject_id": stem}, + ) + for index, window in enumerate(slice_windows(raw.get_data(), WINDOW_SAMPLES)) + ] + + +def main(): + """Slice Sleep-EDFx into 30-second pre-training windows.""" + args = build_arg_parser("Sleep-EDFx pre-training corpus to LMDB").parse_args() + mne.set_log_level("ERROR") + + tasks = [ + (args.input_dir, name) + for name in list_files(args.input_dir, [".edf"]) + ] + write_pretraining_corpus( + "SleepEDFx", tasks, process_recording, args.output_dir, args.num_workers, args.dry_run + ) + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_stew.py b/make_datasets/make_stew.py new file mode 100644 index 0000000..fb0c013 --- /dev/null +++ b/make_datasets/make_stew.py @@ -0,0 +1,167 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the STEW workload corpus into train/val/test LMDBs.""" + +import os + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + build_arg_parser, + referential_coordinates, + run_jobs, +) + +SOURCE_FREQ = 128 +WINDOW_SECONDS = 4 +OVERLAP_SECONDS = 2 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +HOP_SAMPLES = SAMPLING_FREQ * (WINDOW_SECONDS - OVERLAP_SECONDS) +HIGHPASS_FREQ = 1.0 +ASR_CUTOFF = 20 +NUM_SUBJECTS = 48 + +CHANNELS = ["AF3", "F7", "F3", "FC5", "T7", "P7", "O1", "O2", "P8", "T8", "FC6", "F4", "F8", "AF4"] +NUM_CHANNELS = len(CHANNELS) +CHANNEL_COORDS = referential_coordinates(CHANNELS, "AR") + +SUBJECT_SPLITS = { + "train": [1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 23, 25, + 26, 27, 28, 29, 30, 31, 32, 33, 34, 36, 37, 39, 43, 44], + "val": [35, 38, 45, 46], + "test": [22, 40, 41, 47, 48], +} + + +def load_ratings(path: str) -> dict: + """Read the per-subject workload ratings for the low and high workload sessions.""" + ratings = {} + with open(path, "r") as handle: + for line in handle: + parts = [part.strip() for part in line.strip().split(",")] + if len(parts) < 3 or not parts[0].isdigit(): + continue + ratings[int(parts[0])] = { + "lo": None if parts[1] == "" else int(parts[1]), + "hi": None if parts[2] == "" else int(parts[2]), + } + return ratings + + +def rating_to_label(rating): + """Map a 1-9 workload rating onto three balanced classes.""" + if rating is None: + return None + if 1 <= rating <= 3: + return 0 + if 4 <= rating <= 6: + return 1 + if 7 <= rating <= 9: + return 2 + return None + + +def clean_recording(data: np.ndarray) -> np.ndarray: + """High-pass, artifact-correct with ASR, re-reference and resample one recording.""" + import asrpy + + info = mne.create_info(ch_names=CHANNELS, sfreq=SOURCE_FREQ, ch_types="eeg", verbose=False) + raw = mne.io.RawArray(data, info, verbose=False) + raw.filter(l_freq=HIGHPASS_FREQ, h_freq=None, fir_design="firwin", verbose=False) + + asr = asrpy.ASR(sfreq=raw.info["sfreq"], cutoff=ASR_CUTOFF) + asr.fit(raw) + raw = asr.transform(raw) + + raw.set_eeg_reference(ref_channels="average", projection=False, verbose=False) + if int(round(raw.info["sfreq"])) != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, npad="auto", verbose=False) + return raw.get_data().astype(np.float32) + + +def sliding_windows(data: np.ndarray) -> np.ndarray: + """Cut a recording into overlapping fixed-length windows.""" + length = data.shape[1] + if length < WINDOW_SAMPLES: + return np.empty((0, data.shape[0], WINDOW_SAMPLES), dtype=np.float32) + starts = np.arange(0, length - WINDOW_SAMPLES + 1, HOP_SAMPLES, dtype=int) + return np.stack([data[:, start : start + WINDOW_SAMPLES] for start in starts]).astype(np.float32) + + +def process_recording(task): + """Clean one workload session and cut it into labelled windows.""" + split, path, subject, session, label = task + data = np.loadtxt(path) + if data.ndim != 2 or data.shape[1] != NUM_CHANNELS: + return [] + + cleaned = clean_recording(data.T) + samples = [] + for index, window in enumerate(sliding_windows(cleaned)): + samples.append(( + split, + f"sub{subject:02d}_{session}-{index}".encode(), + { + "eeg": window, + "label": int(label), + "channel_coords": CHANNEL_COORDS, + "subject_id": str(subject), + }, + )) + return samples + + +def main(): + """Split STEW by subject and write one LMDB per split.""" + args = build_arg_parser("STEW workload classification to LMDB").parse_args() + mne.set_log_level("ERROR") + + ratings = load_ratings(os.path.join(args.input_dir, "ratings.txt")) + subject_split = { + subject: split for split, subjects in SUBJECT_SPLITS.items() for subject in subjects + } + + tasks = [] + for subject in range(1, NUM_SUBJECTS + 1): + split = subject_split.get(subject) + if split is None or subject not in ratings: + continue + for session in ("lo", "hi"): + path = os.path.join(args.input_dir, f"sub{subject:02d}_{session}.txt") + label = rating_to_label(ratings[subject][session]) + if label is not None and os.path.isfile(path): + tasks.append((split, path, subject, session, label)) + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + for samples in run_jobs(process_recording, tasks, args.num_workers, "STEW recordings"): + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("STEW") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_tuab.py b/make_datasets/make_tuab.py new file mode 100644 index 0000000..afe62ec --- /dev/null +++ b/make_datasets/make_tuab.py @@ -0,0 +1,164 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the TUAB abnormality-detection corpus into train/val/test LMDBs.""" + +import os + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + LMDBWriter, + bipolar_coordinates, + build_arg_parser, + run_jobs, +) + +WINDOW_SECONDS = 10 +WINDOW_SAMPLES = SAMPLING_FREQ * WINDOW_SECONDS +CHANNEL_FOLDER = "01_tcp_ar" +VALIDATION_FRACTION = 0.2 +SPLIT_SEED = 4523 + +BIPOLAR_PAIRS = [ + ("FP1", "F7"), ("F7", "T3"), ("T3", "T5"), ("T5", "O1"), + ("A1", "T3"), ("T3", "C3"), ("C3", "CZ"), ("FP1", "F3"), + ("F3", "C3"), ("C3", "P3"), ("P3", "O1"), + ("FP2", "F8"), ("F8", "T4"), ("T4", "T6"), ("T6", "O2"), + ("T4", "A2"), ("C4", "T4"), ("CZ", "C4"), ("FP2", "F4"), + ("F4", "C4"), ("C4", "P4"), ("P4", "O2"), +] +NUM_CHANNELS = len(BIPOLAR_PAIRS) +CHANNEL_COORDS = bipolar_coordinates(BIPOLAR_PAIRS) + +EDF_PAIRS = [(f"EEG {a}-REF", f"EEG {b}-REF") for a, b in BIPOLAR_PAIRS] +EDF_CHANNELS = list(dict.fromkeys([name for pair in EDF_PAIRS for name in pair])) + + +def subject_of(filename: str) -> str: + """Return the TUAB subject identifier encoded in a filename.""" + return os.path.basename(filename).split("_")[0] + + +def recording_folder(root: str, split: str, label: int) -> str: + """Return the folder holding recordings for one split and label.""" + return os.path.join( + root, "eval" if split == "test" else "train", "abnormal" if label else "normal", CHANNEL_FOLDER + ) + + +def process_subject(task): + """Build the bipolar montage for one subject and cut it into fixed-length windows.""" + split, root, subject, label = task + folder = recording_folder(root, split, label) + samples, skipped = [], [] + + for filename in sorted(os.listdir(folder)): + if not filename.startswith(f"{subject}_") or not filename.endswith(".edf"): + continue + path = os.path.join(folder, filename) + try: + raw = mne.io.read_raw_edf(path, preload=False, verbose=False) + missing = [name for name in EDF_CHANNELS if name not in set(raw.info["ch_names"])] + if missing: + skipped.append((filename, f"missing_channels:{len(missing)}")) + continue + + raw = mne.io.read_raw_edf(path, preload=True, verbose=False) + raw.pick(EDF_CHANNELS) + raw.reorder_channels(EDF_CHANNELS) + raw.notch_filter(60, verbose=False) + raw.filter(l_freq=0.3, h_freq=75, verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, n_jobs=1) + + data = raw.get_data(units="uV") + bipolar = np.stack([ + data[EDF_CHANNELS.index(active)] - data[EDF_CHANNELS.index(reference)] + for active, reference in EDF_PAIRS + ]).astype(np.float32) + + num_windows = bipolar.shape[1] // WINDOW_SAMPLES + if num_windows == 0: + skipped.append((filename, "too_short")) + continue + + for index in range(num_windows): + window = bipolar[:, index * WINDOW_SAMPLES : (index + 1) * WINDOW_SAMPLES] + samples.append(( + split, + f"{filename[:-4]}-{index}".encode(), + { + "eeg": window, + "label": int(label), + "channel_coords": CHANNEL_COORDS, + "subject_id": subject, + }, + )) + except Exception as error: + skipped.append((filename, f"exception:{type(error).__name__}")) + + return samples, skipped + + +def subject_ids(folder: str) -> list: + """Return the sorted unique subject identifiers in a recording folder.""" + return sorted({subject_of(name) for name in os.listdir(folder) if name.endswith(".edf")}) + + +def main(): + """Split TUAB by subject and write one LMDB per split. + + The official evaluation set becomes the test split. Training subjects are divided + into train and validation by subject, so no subject appears in two splits. + """ + args = build_arg_parser("TUAB abnormality detection to LMDB").parse_args() + generator = np.random.default_rng(SPLIT_SEED) + + tasks = [] + for label in (0, 1): + development = subject_ids(recording_folder(args.input_dir, "train", label)) + generator.shuffle(development) + cut = int((1.0 - VALIDATION_FRACTION) * len(development)) + tasks += [("train", args.input_dir, subject, label) for subject in development[:cut]] + tasks += [("val", args.input_dir, subject, label) for subject in development[cut:]] + tasks += [ + ("test", args.input_dir, subject, label) + for subject in subject_ids(recording_folder(args.input_dir, "test", label)) + ] + + writer = LMDBWriter(args.output_dir, dry_run=args.dry_run) + skipped = [] + for samples, reasons in run_jobs(process_subject, tasks, args.num_workers, "TUAB subjects"): + skipped += reasons + for split, key, sample in samples: + writer.put(split, key, sample) + + writer.close() + writer.summarise("TUAB") + if skipped: + print(f" skipped {len(skipped)} recordings, first 5: {skipped[:5]}") + + +if __name__ == "__main__": + main() diff --git a/make_datasets/make_tueg.py b/make_datasets/make_tueg.py new file mode 100644 index 0000000..4c01543 --- /dev/null +++ b/make_datasets/make_tueg.py @@ -0,0 +1,202 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +"""Preprocess the TUEG pre-training corpus into a packed LMDB. + +TUEG is the largest corpus, so each window is stored as a fixed-size byte blob rather +than a pickled dictionary, and the sample keys are written to a companion text file so +readers do not have to scan the database at start-up. + +Subjects that also appear in TUAB are excluded, so that pre-training cannot see any +recording from a subject used in the TUAB fine-tuning splits. +""" + +import os +from pathlib import Path + +import mne +import numpy as np + +from make_datasets.common import ( + SAMPLING_FREQ, + PackedLMDBWriter, + build_arg_parser, + electrode_coordinate, + run_jobs, +) + +SLICE_SECONDS = 30 +WINDOW_SAMPLES = SAMPLING_FREQ * SLICE_SECONDS +MAX_CHANNELS = 64 + +STANDARD_CHANNELS = [ + "EEG FP1-REF", "EEG FP2-REF", "EEG F3-REF", "EEG F4-REF", "EEG C3-REF", "EEG C4-REF", + "EEG P3-REF", "EEG P4-REF", "EEG O1-REF", "EEG O2-REF", "EEG F7-REF", "EEG F8-REF", + "EEG T3-REF", "EEG T4-REF", "EEG T5-REF", "EEG T6-REF", "EEG A1-REF", "EEG A2-REF", + "EEG FZ-REF", "EEG CZ-REF", "EEG PZ-REF", "EEG T1-REF", "EEG T2-REF", +] + +ADDITIONAL_CHANNELS = [ + "EEG FP2-LE", "EEG P3-LE", "EEG C4P-REF", "EEG F8-LE", "EEG P4-LE", "EEG C3-LE", + "EEG T6-LE", "EEG A1-LE", "EEG F4-LE", "EEG T5-LE", "EEG FZ-LE", "EEG O1-LE", + "EEG PZ-LE", "EEG C4-LE", "EEG A2-LE", "EEG CZ-LE", "EEG T1-LE", "EEG F3-LE", + "EEG T2-LE", "EEG FP1-LE", "EEG C3P-REF", "EEG O2-LE", "EEG OZ-LE", "EEG OZ-REF", + "EEG T4-LE", "EEG F7-LE", "EEG SP1-LE", "EEG SP2-LE", "EEG T3-LE", "EEG SP1-REF", + "EEG SP2-REF", +] + +ALL_CHANNELS = list(dict.fromkeys(STANDARD_CHANNELS + ADDITIONAL_CHANNELS)) +ELECTRODE_ALIASES = {"C4P": "CP4", "C3P": "CP3"} + + +def channel_coordinates(channel_name: str, session_reference: str): + """Resolve the two electrode coordinates of one TUEG channel name. + + TUEG mixes average-reference and linked-ear recordings, and the suffix in the + channel name only says which convention the file uses, so the session's reference + is supplied separately. + """ + stripped = channel_name[4:] if channel_name.startswith("EEG ") else channel_name + active, reference = stripped.split("-") + if reference.upper() == "REF": + reference = "AR" + elif reference.upper() == "LE": + reference = session_reference + + def lookup(name: str): + upper = name.upper() + return electrode_coordinate(ELECTRODE_ALIASES.get(upper, upper)) + + return lookup(active), lookup(reference) + + +def process_recording(task): + """Filter, normalise and slice one recording, padding channels to the maximum. + + Each window is min-max normalised per channel to ``[-1, 1]`` before packing, + because the packed format stores no per-sample statistics for a reader to apply. + """ + path, tuab_subjects = task + session = Path(path).stem + subject = session.split("_s")[0] + if subject in tuab_subjects: + return [] + + session_reference = "LE" if ("02_tcp_le" in str(path) or "04_tcp_le" in str(path)) else "AR" + + raw = mne.io.read_raw_edf(path, preload=True, verbose=False) + if raw.n_times == 0: + return [] + + raw.drop_channels([name for name in raw.ch_names if name not in ALL_CHANNELS]) + if not raw.ch_names: + return [] + raw.reorder_channels([name for name in ALL_CHANNELS if name in raw.ch_names]) + + coords = np.asarray( + [channel_coordinates(name, session_reference) for name in raw.ch_names], dtype=np.float32 + ) + + raw.notch_filter(60, verbose=False) + raw.filter(l_freq=0.3, h_freq=75.0, verbose=False) + if raw.info["sfreq"] != SAMPLING_FREQ: + raw.resample(SAMPLING_FREQ, n_jobs=1) + + data = raw.get_data().astype(np.float32) * 1e6 + num_channels, num_timesteps = data.shape + if num_channels > MAX_CHANNELS or num_timesteps < WINDOW_SAMPLES: + return [] + + padded_coords = np.zeros((MAX_CHANNELS, 2, 3), dtype=np.float32) + padded_coords[:num_channels] = coords + + windows = [] + for index in range(num_timesteps // WINDOW_SAMPLES): + window = data[:, index * WINDOW_SAMPLES : (index + 1) * WINDOW_SAMPLES] + minimum = window.min(axis=1, keepdims=True) + maximum = window.max(axis=1, keepdims=True) + normalised = ((window - minimum) / (maximum - minimum + 1e-10) - 0.5) * 2.0 + + padded = np.zeros((MAX_CHANNELS, WINDOW_SAMPLES), dtype=np.float32) + padded[:num_channels] = normalised + windows.append((f"{session}_slice_{index:04d}", padded, padded_coords)) + + return windows + + +def tuab_subject_ids(tuab_dir: str) -> set: + """Return every subject identifier present in the raw TUAB corpus.""" + return {path.name.split("_s")[0] for path in Path(tuab_dir).rglob("*.edf")} + + +def main(): + """Slice TUEG into 30-second pre-training windows, excluding TUAB subjects.""" + parser = build_arg_parser("TUEG pre-training corpus to packed LMDB") + parser.add_argument( + "--tuab_dir", default=None, + help="Root of the raw TUAB corpus. Subjects found here are excluded so they cannot " + "leak into the TUAB fine-tuning splits.", + ) + parser.add_argument( + "--allow_tuab_overlap", action="store_true", + help="Pre-train on all TUEG subjects, including those in TUAB. Only use this when " + "TUAB is not among the downstream evaluation datasets.", + ) + args = parser.parse_args() + mne.set_log_level("ERROR") + + if args.allow_tuab_overlap: + tuab_subjects = set() + print("TUAB exclusion disabled; TUAB subjects may appear in pre-training") + else: + if args.tuab_dir is None: + parser.error( + "--tuab_dir is required so TUAB subjects can be held out of pre-training; " + "pass --allow_tuab_overlap to opt out explicitly" + ) + if not os.path.isdir(args.tuab_dir): + parser.error(f"--tuab_dir does not exist: {args.tuab_dir}") + tuab_subjects = tuab_subject_ids(args.tuab_dir) + if not tuab_subjects: + parser.error( + f"No .edf files found under --tuab_dir {args.tuab_dir}; refusing to continue " + "because this would silently leak TUAB into pre-training" + ) + print(f"Excluding {len(tuab_subjects)} TUAB subjects from pre-training") + + recordings = sorted(str(path) for path in Path(args.input_dir).rglob("*.edf")) + tasks = [(path, tuab_subjects) for path in recordings] + + writer = PackedLMDBWriter( + lmdb_path=os.path.join(args.output_dir, "TUEG.lmdb"), + keys_path=os.path.join(args.output_dir, "TUEG_keys.txt"), + dry_run=args.dry_run, + ) + for windows in run_jobs(process_recording, tasks, args.num_workers, "TUEG recordings"): + for key, waveform, coords in windows: + writer.put(key, waveform, coords) + + writer.close() + writer.summarise("TUEG") + + +if __name__ == "__main__": + main() diff --git a/models/README.md b/models/README.md index 81c5673..0951900 100644 --- a/models/README.md +++ b/models/README.md @@ -4,6 +4,8 @@ Copyright (C) 2025-2026 ETH Zurich, Switzerland. SPDX-License-Identifier: Apache This directory contains the PyTorch `nn.Module` implementations for the BioFoundation model families. Hydra model settings live in [`../config/model`](../config/model/), while the canonical family metadata and batch requirements live in [`../biofoundation/model_registry.py`](../biofoundation/model_registry.py). +Families come in two shapes. The five original families bundle their output layer into the model and select it from `num_classes` at construction time. S-CEReBrO separates the two: [`s_cerebro.py`](s_cerebro.py) is an encoder that emits token embeddings, and [`model_heads`](model_heads/) holds the prediction heads that consume them, configured through [`../config/model_head`](../config/model_head/). Both shapes are supported; the [protocols](../biofoundation/core/protocols.py) describe the second. + ## Available Models | Model | Signals | Summary | Resources | @@ -13,5 +15,6 @@ This directory contains the PyTorch `nn.Module` implementations for the BioFound | TinyMyo | sEMG | Compact rotary Transformer designed for flexible EMG processing and edge deployment. | [Documentation](../docs/model/TinyMyo.md) / [Hugging Face](https://huggingface.co/PulpBio/TinyMyo) | | LuMamba | EEG | LUNA-style channel unification with efficient Mamba temporal modeling. | [Documentation](../docs/model/LuMamba.md) / [Hugging Face](https://huggingface.co/PulpBio/LuMamba) | | PanLUNA | EEG, ECG, PPG | Sensor-aware query unification for unimodal and multimodal biosignal learning. | [Documentation](../docs/model/PanLUNA.md) / [Hugging Face](https://huggingface.co/PulpBio/PanLUNA) | +| S-CEReBrO | EEG | Windowed alternating attention over per-channel patches, with a separate prediction head. | [Documentation](../docs/model/SCEReBrO.md) / [Hugging Face](https://huggingface.co/PulpBio/S-CEReBrO) | Use the matching pre-training or fine-tuning experiment in [`../config/experiment`](../config/experiment/) rather than instantiating a model in isolation when starting a reproducible run. diff --git a/models/model_heads/__init__.py b/models/model_heads/__init__.py new file mode 100644 index 0000000..3cc5778 --- /dev/null +++ b/models/model_heads/__init__.py @@ -0,0 +1,20 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* diff --git a/models/model_heads/mlp_classification_head.py b/models/model_heads/mlp_classification_head.py new file mode 100644 index 0000000..4c4eaa0 --- /dev/null +++ b/models/model_heads/mlp_classification_head.py @@ -0,0 +1,118 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn +from timm.layers import trunc_normal_ + + +class MlpClassificationHead(nn.Module): + """Classification head over the encoder's token embeddings. + + An optional token-wise pooler is applied first. Tokens are then aggregated + either by averaging (``pooling_method='mean'``) or by concatenation + (``pooling_method='flatten'``), and a linear or MLP classifier produces one + prediction per input window. + + Args: + embed_dim: Token embedding dimension. + num_classes: Number of output classes. + pooling_method: ``'mean'`` or ``'flatten'``. + dropout: Dropout used in the pooler and, for ``'flatten'``, the classifier. + pooling: Whether to apply the token-wise pooler. + num_channels: Channel count, required for ``'flatten'`` to size the classifier. + num_patches: Patches per channel, required for ``'flatten'``. + """ + + def __init__( + self, + embed_dim: int = 200, + num_classes: int = 2, + pooling_method: str = "mean", + dropout: float = 0.1, + pooling: bool = True, + num_channels: int = 32, + num_patches: int = 10, + ): + super().__init__() + if pooling_method not in {"mean", "flatten"}: + raise ValueError(f"pooling_method must be 'mean' or 'flatten', got '{pooling_method}'") + + self.embed_dim = embed_dim + self.num_classes = num_classes + self.pooling_method = pooling_method + self.pooling = pooling + self.num_tokens = num_channels * num_patches + + if self.pooling: + self.pooler = nn.Sequential( + nn.Linear(embed_dim, embed_dim), + nn.GELU(), + nn.LayerNorm(embed_dim), + nn.Dropout(dropout), + nn.Tanh(), + ) + + if pooling_method == "mean": + self.classifier = nn.Linear(embed_dim, num_classes) + else: + self.classifier = nn.Sequential( + nn.Linear(self.num_tokens * embed_dim, num_patches * embed_dim), + nn.ELU(), + nn.Dropout(dropout), + nn.Linear(num_patches * embed_dim, embed_dim), + nn.ELU(), + nn.Dropout(dropout), + nn.Linear(embed_dim, num_classes), + ) + + self.apply(self._init_weights) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Classify a batch of token sequences. + + Args: + x: Token embeddings of shape ``(batch, num_tokens, embed_dim)``. + + Returns: + Logits of shape ``(batch, num_classes)``. + """ + if self.pooling: + x = self.pooler(x) + + if self.pooling_method == "mean": + x = x.mean(dim=1) + else: + batch, num_tokens, embed_dim = x.shape + if num_tokens != self.num_tokens or embed_dim != self.embed_dim: + raise ValueError( + f"Expected ({self.num_tokens}, {self.embed_dim}) tokens, got ({num_tokens}, {embed_dim})" + ) + x = x.reshape(batch, num_tokens * embed_dim) + + return self.classifier(x) diff --git a/models/model_heads/mlp_regression_head.py b/models/model_heads/mlp_regression_head.py new file mode 100644 index 0000000..0bfc211 --- /dev/null +++ b/models/model_heads/mlp_regression_head.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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn +from timm.layers import trunc_normal_ + + +class MlpRegressionHead(nn.Module): + """Scalar regression head over the encoder's token embeddings. + + An optional token-wise pooler is applied, tokens are mean-pooled, and a linear + layer produces one scalar per input window. When ``bounded_output`` is set the + prediction passes through a sigmoid, which suits targets defined on ``[0, 1]`` + such as SEED-VIG PERCLOS; disable it for unbounded targets. + """ + + def __init__( + self, + embed_dim: int = 200, + dropout: float = 0.1, + pooling: bool = True, + bounded_output: bool = True, + ): + super().__init__() + self.embed_dim = embed_dim + self.pooling = pooling + self.bounded_output = bounded_output + + if self.pooling: + self.pooler = nn.Sequential( + nn.Linear(embed_dim, embed_dim), + nn.GELU(), + nn.LayerNorm(embed_dim), + nn.Dropout(dropout), + nn.Tanh(), + ) + + layers = [nn.Linear(embed_dim, 1)] + if bounded_output: + layers.append(nn.Sigmoid()) + self.regressor = nn.Sequential(*layers) + + self.apply(self._init_weights) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Predict one scalar per input window. + + Args: + x: Token embeddings of shape ``(batch, num_tokens, embed_dim)``. + + Returns: + Predictions of shape ``(batch,)``. + """ + if self.pooling: + x = self.pooler(x) + + return self.regressor(x.mean(dim=1)).squeeze(-1) diff --git a/models/model_heads/patch_reconstruction_head.py b/models/model_heads/patch_reconstruction_head.py new file mode 100644 index 0000000..d654852 --- /dev/null +++ b/models/model_heads/patch_reconstruction_head.py @@ -0,0 +1,48 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn + + +class PatchReconstructionHead(nn.Module): + """Linear decoder mapping each token embedding back to its waveform patch. + + This is the SimMIM-style decoder used for pre-training: every token is projected + independently, with no decoder transformer and no re-ordering of the sequence. + """ + + def __init__(self, embed_dim: int = 200, patch_size: int = 200): + super().__init__() + self.embed_dim = embed_dim + self.patch_size = patch_size + self.decoder_pred = nn.Linear(embed_dim, patch_size, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Reconstruct waveform patches. + + Args: + x: Token embeddings of shape ``(batch, num_tokens, embed_dim)``. + + Returns: + Reconstructed patches of shape ``(batch, num_tokens, patch_size)``. + """ + return self.decoder_pred(x) diff --git a/models/model_heads/sequence_classification_head.py b/models/model_heads/sequence_classification_head.py new file mode 100644 index 0000000..56ad3eb --- /dev/null +++ b/models/model_heads/sequence_classification_head.py @@ -0,0 +1,130 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn +import torch.nn.functional as F +from timm.layers import trunc_normal_ + + +class SequenceClassificationHead(nn.Module): + """Per-epoch classification head for sequences of consecutive EEG epochs. + + Used for sleep staging (ISRUC), where a sample is a sequence of + ``sequence_length`` consecutive epochs and every epoch carries its own label. + The encoder sees the sequence flattened into the batch axis, so this head + receives ``(batch * sequence_length, tokens_per_epoch, embed_dim)``. It regroups + the sequence, compresses each epoch to a single vector, contextualises the epochs + against one another with a small transformer encoder, and emits one prediction per + epoch. + + Logits are returned flattened as ``(batch * sequence_length, num_classes)`` so + they align with the labels flattened by the classification task. + + Args: + sequence_length: Number of consecutive epochs per sample. + num_channels: EEG channels per epoch. + num_patches: Patches per channel per epoch. + embed_dim: Token embedding dimension of the encoder. + num_classes: Number of classes predicted per epoch. + hidden_dim: Width of the per-epoch representation and the sequence encoder. + num_layers: Transformer encoder layers over the epoch sequence. + nhead: Attention heads in the sequence encoder. + dim_feedforward: Feed-forward width in the sequence encoder. + dropout: Dropout in the sequence encoder. + norm_first: Use pre-norm ordering in the sequence encoder. + """ + + def __init__( + self, + sequence_length: int, + num_channels: int, + num_patches: int, + embed_dim: int, + num_classes: int, + hidden_dim: int = 512, + num_layers: int = 1, + nhead: int = 4, + dim_feedforward: int = 2048, + dropout: float = 0.1, + norm_first: bool = True, + ): + super().__init__() + self.sequence_length = sequence_length + self.num_channels = num_channels + self.num_patches = num_patches + self.embed_dim = embed_dim + self.num_classes = num_classes + self.tokens_per_epoch = num_channels * num_patches + self.feature_dim = self.tokens_per_epoch * embed_dim + + self.head = nn.Sequential( + nn.Linear(self.feature_dim, hidden_dim), + nn.GELU(), + ) + + encoder_layer = nn.TransformerEncoderLayer( + d_model=hidden_dim, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=F.gelu, + batch_first=True, + norm_first=norm_first, + ) + self.sequence_encoder = nn.TransformerEncoder( + encoder_layer, num_layers=num_layers, enable_nested_tensor=False + ) + self.classifier = nn.Linear(hidden_dim, num_classes) + + self.apply(self._init_weights) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Classify every epoch in each sequence. + + Args: + x: Token embeddings of shape ``(batch * sequence_length, tokens_per_epoch, embed_dim)``. + + Returns: + Logits of shape ``(batch * sequence_length, num_classes)``. + """ + flat_batch, num_tokens, embed_dim = x.shape + if embed_dim != self.embed_dim: + raise ValueError(f"Expected embed_dim {self.embed_dim}, got {embed_dim}") + if num_tokens != self.tokens_per_epoch: + raise ValueError(f"Expected {self.tokens_per_epoch} tokens per epoch, got {num_tokens}") + if flat_batch % self.sequence_length != 0: + raise ValueError( + f"Flattened batch {flat_batch} is not divisible by sequence_length {self.sequence_length}" + ) + + batch = flat_batch // self.sequence_length + x = x.reshape(batch, self.sequence_length, self.feature_dim) + x = self.head(x) + x = self.sequence_encoder(x) + return self.classifier(x).reshape(flat_batch, self.num_classes) diff --git a/models/modules/attention.py b/models/modules/attention.py new file mode 100644 index 0000000..0b9e21d --- /dev/null +++ b/models/modules/attention.py @@ -0,0 +1,560 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +from typing import Dict, Iterable, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from timm.layers import DropPath, Mlp + + +def build_window_indices( + length: int, + window_size: int, + dilation: int, + include_self: bool, + shift: int, + device: torch.device, +) -> torch.Tensor: + """Build clamped, dilated, shifted window indices for every position along an axis. + + Args: + length: Number of positions along the axis. + window_size: Requested window size; clamped to ``length``. + dilation: Spacing between consecutive window offsets. + include_self: Whether a position attends to itself. + shift: Constant offset applied to the whole window. + device: Device of the returned index tensor. + + Returns: + Long tensor of shape ``(length, effective_window_size)``. + """ + step = max(int(dilation), 1) + effective_size = max(1, min(int(window_size), int(length))) + offset = int(shift) + + if include_self: + left = (effective_size - 1) // 2 + right = effective_size - 1 - left + offsets = torch.cat([ + torch.arange(-left, 0, device=device), + torch.zeros(1, device=device, dtype=torch.long), + torch.arange(1, right + 1, device=device), + ]) * step + offset + else: + left = effective_size // 2 + right = effective_size - left + offsets = torch.cat([ + torch.arange(-left, 0, device=device), + torch.arange(1, right + 1, device=device), + ]) * step + offset + + base = torch.arange(length, device=device)[:, None] + return (base + offsets[None, :]).clamp_(0, length - 1).to(torch.long) + + +class WindowedAlternatingAttention(nn.Module): + """Windowed attention that alternates between the spatial and temporal axes. + + Tokens are laid out as ``num_channels * num_patches`` per sample. Even-indexed + blocks attend across channels at a fixed time index (spatial pass); odd-indexed + blocks attend across time within a fixed channel (temporal pass). Each pass is + restricted to a dilated, optionally shifted window, so cost per block is + ``O(batch * channels * patches * window_size)`` instead of quadratic in the + full token count. + + Setting ``use_axial_mode`` replaces the alternating schedule with an axial one: + the first half of the blocks perform spatial attention and the second half + temporal attention. + + Padded keys are masked to ``-inf`` before the softmax and padded queries are + zeroed on output, so variable channel counts share one batch safely. + """ + + def __init__( + self, + dim: int, + num_heads: int = 8, + *, + num_channels: int, + block_idx: int = 0, + total_blocks: Optional[int] = None, + use_axial_mode: bool = False, + spatial_first: bool = True, + window_size_spatial: int = 7, + window_size_temporal: int = 5, + dilation_spatial: int = 1, + dilation_temporal: int = 1, + include_self: bool = True, + qkv_bias: bool = False, + qk_norm: bool = True, + normalize_qk: bool = False, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + + self.num_channels = int(num_channels) + self.block_idx = int(block_idx) + self.use_axial_mode = bool(use_axial_mode) + + if self.use_axial_mode: + if total_blocks is None: + raise ValueError("total_blocks is required when use_axial_mode is True") + half = total_blocks // 2 + self.spatial_pass = block_idx < half if spatial_first else block_idx >= half + else: + self.spatial_pass = (block_idx % 2 == 0) + + self.window_size_spatial = int(window_size_spatial) + self.window_size_temporal = int(window_size_temporal) + self.dilation_spatial = int(dilation_spatial) + self.dilation_temporal = int(dilation_temporal) + self.include_self = bool(include_self) + self.normalize_qk = bool(normalize_qk) + + self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias) + self.q_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = nn.LayerNorm(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + self._index_cache: Dict[Tuple[str, int, int], torch.Tensor] = {} + + def _window_indices( + self, axis: str, length: int, window_size: int, dilation: int, shift: int, device: torch.device + ) -> torch.Tensor: + """Return cached window indices for one axis, rebuilding on device change.""" + key = (axis, length, shift) + cached = self._index_cache.get(key) + if cached is None or cached.device != device: + cached = build_window_indices( + length=length, + window_size=window_size, + dilation=dilation, + include_self=self.include_self, + shift=shift, + device=device, + ) + self._index_cache[key] = cached + return cached + + def _windowed_softmax_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_mask: Optional[torch.Tensor], + indices: torch.Tensor, + ) -> torch.Tensor: + """Attend from each position to its window along dim 1. + + Args: + q, k, v: Tensors of shape ``(flat_batch, length, heads, head_dim)``. + key_mask: Optional ``(flat_batch, length)`` mask, 1 for real and 0 for padded keys. + indices: ``(length, window)`` key indices per query. + + Returns: + Tensor of shape ``(flat_batch, length, heads, head_dim)``. + """ + length, window = indices.shape + gather = indices.view(1, length, window, 1, 1) + + keys = torch.take_along_dim(k.unsqueeze(2), gather, dim=1) + values = torch.take_along_dim(v.unsqueeze(2), gather, dim=1) + + logits = (q.unsqueeze(2) * keys).sum(-1) * self.scale + + fully_masked = None + if key_mask is not None: + mask = key_mask.to(logits.dtype) + mask_window = torch.take_along_dim(mask.unsqueeze(2), indices.view(1, length, window), dim=1) + logits = logits.masked_fill(mask_window.unsqueeze(-1) <= 0, float("-inf")) + fully_masked = (mask_window.sum(dim=2, keepdim=True) == 0) + logits = torch.where(fully_masked.unsqueeze(-1), torch.zeros_like(logits), logits) + + attn = torch.softmax(logits, dim=2).to(values.dtype) + if fully_masked is not None: + attn = torch.where(fully_masked.unsqueeze(-1), torch.zeros_like(attn), attn) + attn = self.attn_drop(attn) + + return (attn.unsqueeze(-1) * values).sum(dim=2) + + def forward( + self, + x: torch.Tensor, + attn_mask: Optional[torch.Tensor] = None, + *, + shift_spatial: int = 0, + shift_temporal: int = 0, + ) -> torch.Tensor: + """Apply one windowed spatial or temporal attention pass. + + Args: + x: Token embeddings of shape ``(batch, num_channels * num_patches, dim)``. + attn_mask: Optional ``(batch, num_tokens)`` mask, 1 for real and 0 for padded tokens. + shift_spatial: Window shift applied on a spatial pass. + shift_temporal: Window shift applied on a temporal pass. + + Returns: + Tensor of shape ``(batch, num_tokens, dim)``. + """ + batch, num_tokens, dim = x.shape + channels = self.num_channels + if num_tokens % channels != 0: + raise ValueError("num_tokens must be divisible by num_channels") + patches = num_tokens // channels + + qkv = self.qkv(x).reshape(batch, num_tokens, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + if self.normalize_qk: + q = F.normalize(q, dim=-1) + k = F.normalize(k, dim=-1) + + shape = (batch, channels, patches, self.num_heads, self.head_dim) + q = q.permute(0, 2, 1, 3).reshape(shape) + k = k.permute(0, 2, 1, 3).reshape(shape) + v = v.permute(0, 2, 1, 3).reshape(shape) + mask = attn_mask.view(batch, channels, patches) if attn_mask is not None else None + + if self.spatial_pass: + indices = self._window_indices( + "spatial", channels, self.window_size_spatial, self.dilation_spatial, shift_spatial, x.device + ) + flat = (batch * patches, channels, self.num_heads, self.head_dim) + q_flat = q.permute(0, 2, 1, 3, 4).reshape(flat) + k_flat = k.permute(0, 2, 1, 3, 4).reshape(flat) + v_flat = v.permute(0, 2, 1, 3, 4).reshape(flat) + mask_flat = mask.permute(0, 2, 1).reshape(batch * patches, channels) if mask is not None else None + out = self._windowed_softmax_attention(q_flat, k_flat, v_flat, mask_flat, indices) + out = out.reshape(batch, patches, channels, self.num_heads, self.head_dim).permute(0, 2, 1, 3, 4) + else: + indices = self._window_indices( + "temporal", patches, self.window_size_temporal, self.dilation_temporal, shift_temporal, x.device + ) + flat = (batch * channels, patches, self.num_heads, self.head_dim) + q_flat = q.reshape(flat) + k_flat = k.reshape(flat) + v_flat = v.reshape(flat) + mask_flat = mask.reshape(batch * channels, patches) if mask is not None else None + out = self._windowed_softmax_attention(q_flat, k_flat, v_flat, mask_flat, indices) + out = out.reshape(batch, channels, patches, self.num_heads, self.head_dim) + + out = out.reshape(batch, num_tokens, dim) + out = self.proj_drop(self.proj(out)) + + if attn_mask is not None: + out = out * attn_mask.unsqueeze(-1).to(out.dtype) + + return out + + +class AlternatingAttention(nn.Module): + """Full attention that alternates between the spatial and temporal axes. + + Even-indexed blocks attend over all channels at a fixed time index, odd-indexed + blocks over all time indices within a channel. This is the unwindowed ablation + of :class:`WindowedAlternatingAttention`. + """ + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_norm: bool = False, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + num_channels: int = 64, + block_idx: int = 0, + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + self.num_channels = num_channels + self.spatial_pass = (block_idx % 2 == 0) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Apply one full spatial or temporal attention pass. + + Args: + x: Token embeddings of shape ``(batch, num_channels * num_patches, dim)``. + attn_mask: Optional ``(batch, num_tokens)`` mask, 1 for real and 0 for padded tokens. + + Returns: + Tensor of shape ``(batch, num_tokens, dim)``. + """ + num_tokens = x.shape[1] + patches = num_tokens // self.num_channels + + if self.spatial_pass: + x = rearrange(x, "B (C T) D -> (B T) C D", C=self.num_channels) + if attn_mask is not None: + attn_mask = rearrange(attn_mask, "B (C T) -> (B T) C", C=self.num_channels) + x = self._attend(x, attn_mask) + return rearrange(x, "(B T) C D -> B (C T) D", T=patches) + + x = rearrange(x, "B (C T) D -> (B C) T D", C=self.num_channels) + if attn_mask is not None: + attn_mask = rearrange(attn_mask, "B (C T) -> (B C) T", C=self.num_channels) + x = self._attend(x, attn_mask) + return rearrange(x, "(B C) T D -> B (C T) D", C=self.num_channels) + + def _attend(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Scaled dot-product attention over the full second axis of ``x``.""" + batch, length, dim = x.shape + qkv = self.qkv(x).reshape(batch, length, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + attn = (q * self.scale) @ k.transpose(-2, -1) + if attn_mask is not None: + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1).expand(batch, self.num_heads, length, length) + attn = attn.masked_fill(attn_mask == 0, float("-inf")) + + attn = attn.softmax(dim=-1) + if attn_mask is not None: + attn = attn.masked_fill(attn_mask.sum(dim=-1, keepdim=True).eq(0), 0.0) + + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2).reshape(batch, length, dim) + return self.proj_drop(self.proj(x)) + + +class FullAttention(nn.Module): + """Standard multi-head self-attention over the complete token sequence.""" + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_norm: bool = False, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + norm_layer: nn.Module = nn.LayerNorm, + ) -> None: + super().__init__() + if dim % num_heads != 0: + raise ValueError("dim must be divisible by num_heads") + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity() + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Attend over all tokens. + + Args: + x: Token embeddings of shape ``(batch, num_tokens, dim)``. + attn_mask: Optional ``(batch, num_tokens)`` mask, 1 for real and 0 for padded tokens. + + Returns: + Tensor of shape ``(batch, num_tokens, dim)``. + """ + batch, num_tokens, dim = x.shape + qkv = self.qkv(x).reshape(batch, num_tokens, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + q, k = self.q_norm(q), self.k_norm(k) + + attn = (q * self.scale) @ k.transpose(-2, -1) + if attn_mask is not None: + expanded = attn_mask.unsqueeze(1).unsqueeze(1).expand(batch, self.num_heads, num_tokens, num_tokens) + attn = attn.masked_fill(expanded == 0, float("-inf")) + + attn = attn.softmax(dim=-1) + if attn_mask is not None: + attn = attn.masked_fill(expanded.sum(dim=-1, keepdim=True).eq(0), 0.0) + + attn = self.attn_drop(attn) + x = (attn @ v).transpose(1, 2).reshape(batch, num_tokens, dim) + return self.proj_drop(self.proj(x)) + + +class LayerScale(nn.Module): + """Per-channel learnable rescaling of a residual branch.""" + + def __init__(self, dim: int, init_values: float = 1e-5, inplace: bool = False) -> None: + super().__init__() + self.inplace = inplace + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.mul_(self.gamma) if self.inplace else x * self.gamma + + +class TransformerBlock(nn.Module): + """Pre-norm transformer block with a configurable attention mechanism. + + ``attention_type`` selects one of ``windowed-alternating`` (the CEReBrO + default), ``alternating``, or ``full``. For the windowed variant, the + dilation and shift schedules are indexed by spatial/temporal *pair* rather + than by block, so a spatial block and the temporal block that follows it + share the same schedule entry. + """ + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + qk_norm: bool = False, + proj_drop: float = 0.0, + attn_drop: float = 0.0, + init_values: Optional[float] = None, + drop_path: float = 0.0, + act_layer: nn.Module = nn.GELU, + norm_layer: nn.Module = nn.LayerNorm, + mlp_layer: nn.Module = Mlp, + num_channels: int = 23, + attention_type: str = "windowed-alternating", + block_idx: int = 0, + total_blocks: int = 12, + spatial_first: bool = True, + window_size_spatial: int = 7, + window_size_temporal: int = 5, + dilation_cycle_spatial: Iterable[int] = (1, 2, 4), + dilation_cycle_temporal: Iterable[int] = (1, 2, 4), + shift_cycle_spatial: Iterable[int] = (-1, 1, -2, 2), + shift_cycle_temporal: Iterable[int] = (-1, 1, -2, 2), + include_self: bool = True, + normalize_qk: bool = False, + use_axial_mode: bool = False, + ) -> None: + super().__init__() + self.attention_type = attention_type + self.block_idx = block_idx + + dilation_cycle_spatial = tuple(dilation_cycle_spatial) + dilation_cycle_temporal = tuple(dilation_cycle_temporal) + shift_cycle_spatial = tuple(shift_cycle_spatial) + shift_cycle_temporal = tuple(shift_cycle_temporal) + pair_idx = block_idx // 2 + + self.shift_spatial = shift_cycle_spatial[pair_idx % len(shift_cycle_spatial)] + self.shift_temporal = shift_cycle_temporal[pair_idx % len(shift_cycle_temporal)] + + self.norm1 = norm_layer(dim) + + if attention_type == "windowed-alternating": + self.attn = WindowedAlternatingAttention( + dim=dim, + num_heads=num_heads, + num_channels=num_channels, + block_idx=block_idx, + total_blocks=total_blocks, + use_axial_mode=use_axial_mode, + spatial_first=spatial_first, + window_size_spatial=window_size_spatial, + window_size_temporal=window_size_temporal, + dilation_spatial=dilation_cycle_spatial[pair_idx % len(dilation_cycle_spatial)], + dilation_temporal=dilation_cycle_temporal[pair_idx % len(dilation_cycle_temporal)], + include_self=include_self, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + normalize_qk=normalize_qk, + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + elif attention_type == "alternating": + self.attn = AlternatingAttention( + dim, + num_heads=num_heads, + num_channels=num_channels, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + block_idx=block_idx, + ) + elif attention_type == "full": + self.attn = FullAttention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_norm=qk_norm, + attn_drop=attn_drop, + proj_drop=proj_drop, + norm_layer=norm_layer, + ) + else: + raise ValueError( + f"Unknown attention_type '{attention_type}'; " + "expected 'windowed-alternating', 'alternating', or 'full'" + ) + + self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = mlp_layer( + in_features=dim, + hidden_features=int(dim * mlp_ratio), + act_layer=act_layer, + drop=proj_drop, + ) + self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() + self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """Run attention and the feed-forward network with residual connections.""" + if self.attention_type == "windowed-alternating": + attended = self.attn( + self.norm1(x), + attn_mask, + shift_spatial=self.shift_spatial, + shift_temporal=self.shift_temporal, + ) + else: + attended = self.attn(self.norm1(x), attn_mask) + + x = x + self.drop_path1(self.ls1(attended)) + return x + self.drop_path2(self.ls2(self.mlp(self.norm2(x)))) diff --git a/models/modules/patching.py b/models/modules/patching.py new file mode 100644 index 0000000..207d544 --- /dev/null +++ b/models/modules/patching.py @@ -0,0 +1,137 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn +from einops import rearrange + + +def patchify(signal: torch.Tensor, patch_size: int) -> torch.Tensor: + """Split a multi-channel waveform into per-channel patches. + + Args: + signal: Waveform of shape ``(batch, num_channels, num_timesteps)``. + patch_size: Number of timesteps per patch; must divide ``num_timesteps``. + + Returns: + Tensor of shape ``(batch, num_channels * num_patches, patch_size)``, with + channel as the slower-varying index. + """ + batch, channels, timesteps = signal.shape + if timesteps % patch_size != 0: + raise ValueError(f"num_timesteps ({timesteps}) must be divisible by patch_size ({patch_size})") + patches = signal.reshape(batch, channels, timesteps // patch_size, patch_size) + return rearrange(patches, "B C t p -> B (C t) p") + + +def unpatchify(patches: torch.Tensor, num_channels: int) -> torch.Tensor: + """Reassemble per-channel patches into a multi-channel waveform. + + Args: + patches: Tensor of shape ``(batch, num_channels * num_patches, patch_size)``. + num_channels: Number of channels encoded in the token axis. + + Returns: + Waveform of shape ``(batch, num_channels, num_timesteps)``. + """ + return rearrange(patches, "B (C t) p -> B C (t p)", C=num_channels) + + +PATCH_SIZE = 200 +CONV_OUT_CHANNELS = 8 +CONV_OUTPUT_FEATURES = 200 + + +class TemporalConvTokenizer(nn.Module): + """Convolutional tokenizer mapping one waveform patch to one embedding. + + Each ``(channel, patch)`` pair is encoded independently by a stack of three + strided 1D convolutions (applied as 2D convolutions over a flattened + channel-patch axis), then projected to ``embed_dim``. The channel dimension is + preserved, so the token count is ``num_channels * num_patches``. + + The convolution stack is defined for a patch of 200 timesteps, which at the + project-wide 200 Hz sampling rate is one second per token. The first convolution + strides by 8, so a 200-sample patch yields 25 positions of 8 features each and the + projection consumes exactly 200 features. Other patch sizes are rejected rather + than silently reshaped. + """ + + def __init__( + self, + patch_size: int = PATCH_SIZE, + out_channels: int = CONV_OUT_CHANNELS, + embed_dim: int = PATCH_SIZE, + ): + super().__init__() + if patch_size != PATCH_SIZE: + raise ValueError( + f"TemporalConvTokenizer is defined for patch_size={PATCH_SIZE}, got {patch_size}" + ) + self.patch_size = patch_size + self.out_channels = out_channels + self.embed_dim = embed_dim + + self.conv1 = nn.Conv2d(1, out_channels, kernel_size=(1, 15), stride=(1, 8), padding=(0, 7)) + self.gelu1 = nn.GELU() + self.norm1 = nn.GroupNorm(4, out_channels) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)) + self.gelu2 = nn.GELU() + self.norm2 = nn.GroupNorm(4, out_channels) + self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size=(1, 3), padding=(0, 1)) + self.gelu3 = nn.GELU() + self.norm3 = nn.GroupNorm(4, out_channels) + self.proj = nn.Linear(CONV_OUTPUT_FEATURES, embed_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Embed patched waveforms. + + Args: + x: Patched waveform of shape ``(batch, num_channels, num_patches, patch_size)``. + + Returns: + Tensor of shape ``(batch, num_channels * num_patches, embed_dim)``. + """ + x = rearrange(x, "B C P S -> B (C P) S").unsqueeze(1) + x = self.gelu1(self.norm1(self.conv1(x))) + x = self.gelu2(self.norm2(self.conv2(x))) + x = self.gelu3(self.norm3(self.conv3(x))) + x = rearrange(x, "B F N T -> B N (T F)") + return self.proj(x) + + +class PatchEmbedding(nn.Module): + """Waveform patch embedding used by S-CEReBrO. + + Thin wrapper around :class:`TemporalConvTokenizer`. The inner module is held in + an attribute named ``patch_embed`` so that parameter keys remain + ``patch_embed.patch_embed.*``, keeping checkpoints from earlier versions of + this code loadable. Weight initialisation is delegated to the encoder, which + applies it to the whole model tree. + """ + + def __init__(self, patch_size: int = PATCH_SIZE, embed_dim: int = PATCH_SIZE): + super().__init__() + self.patch_embed = TemporalConvTokenizer(patch_size=patch_size, embed_dim=embed_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Embed patched waveforms of shape ``(batch, channels, patches, patch_size)``.""" + return self.patch_embed(x) diff --git a/models/modules/pos_embed.py b/models/modules/pos_embed.py new file mode 100644 index 0000000..9540dec --- /dev/null +++ b/models/modules/pos_embed.py @@ -0,0 +1,90 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +import torch.nn as nn +from timm.layers import trunc_normal_ + + +class PositionalEmbedding(nn.Module): + """Learned temporal position embedding, shared across channels. + + A single table of ``max_patches`` vectors is broadcast over the channel axis, so + two tokens at the same time index in different channels receive the same + temporal embedding. + """ + + def __init__(self, max_patches: int, num_channels: int, embed_dim: int): + super().__init__() + self.max_patches = max_patches + self.num_channels = num_channels + self.embed_dim = embed_dim + self.pos_embedding = nn.Parameter(torch.zeros(1, max_patches, embed_dim)) + trunc_normal_(self.pos_embedding, std=0.02, a=-0.02, b=0.02) + + def forward(self) -> torch.Tensor: + """Return the embedding table of shape ``(1, num_channels, max_patches, embed_dim)``.""" + return self.pos_embedding.unsqueeze(1).repeat(1, self.num_channels, 1, 1) + + +class ChannelEmbedding(nn.Module): + """Channel embedding computed from 3D electrode coordinates. + + Every channel is described by two electrodes (a bipolar pair, or a scalp + electrode and its reference). A shared MLP maps each electrode's 3D coordinate + to ``embed_dim // 2`` features and the two halves are concatenated. Because the + embedding is a function of geometry rather than of a channel index, montages + with different channel counts and orderings share the same parameters. + """ + + def __init__(self, embed_dim: int): + super().__init__() + self.embed_dim = embed_dim + self.mlp = nn.Sequential( + nn.Linear(3, embed_dim // 4), + nn.GELU(), + nn.Linear(embed_dim // 4, embed_dim // 4), + nn.GELU(), + nn.Linear(embed_dim // 4, embed_dim // 2), + ) + for layer in self.mlp: + if isinstance(layer, nn.Linear): + nn.init.kaiming_normal_(layer.weight, mode="fan_out", nonlinearity="relu") + if layer.bias is not None: + nn.init.constant_(layer.bias, 0) + + def forward(self, channel_positions: torch.Tensor) -> torch.Tensor: + """Embed electrode coordinates. + + Args: + channel_positions: Tensor of shape ``(batch, num_channels, 2, 3)``. + + Returns: + Tensor of shape ``(batch, num_channels, embed_dim)``. + """ + batch, channels, electrodes, coords = channel_positions.shape + if electrodes != 2: + raise ValueError(f"Expected 2 electrodes per channel, got {electrodes}") + if coords != 3: + raise ValueError(f"Expected 3D electrode coordinates, got {coords}") + + embedded = self.mlp(channel_positions.reshape(-1, 3)) + return embedded.view(batch, channels, self.embed_dim) diff --git a/models/s_cerebro.py b/models/s_cerebro.py new file mode 100644 index 0000000..9702f0a --- /dev/null +++ b/models/s_cerebro.py @@ -0,0 +1,208 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +from typing import Iterable, Optional + +import torch +import torch.nn as nn +from timm.layers import trunc_normal_ + +from models.modules.attention import TransformerBlock +from models.modules.patching import PatchEmbedding +from models.modules.pos_embed import ChannelEmbedding, PositionalEmbedding + + +class SCerebroEncoder(nn.Module): + """S-CEReBrO: a transformer encoder for multi-channel EEG with windowed alternating attention. + + A waveform of shape ``(batch, num_channels, num_patches, patch_size)`` is + tokenised per channel-patch pair, giving ``num_channels * num_patches`` tokens. + Each token receives a learned temporal position embedding plus a channel + embedding derived from its electrode coordinates. The token sequence then passes + through ``depth`` transformer blocks whose attention alternates between the + channel (spatial) and time (temporal) axes, restricted to dilated and shifted + windows on each axis. + + The temporal position table is sized by ``max_timesteps // patch_size`` and + sliced to the number of patches actually present, so one pre-trained encoder can + be fine-tuned on shorter recordings and on montages with fewer channels than + ``max_channels``. + + Args: + patch_size: Timesteps per patch. + num_channels: Number of EEG channels in the input. + embed_dim: Token embedding dimension. + depth: Number of transformer blocks. + num_heads: Attention heads per block; must divide ``embed_dim``. + mlp_ratio: Feed-forward hidden size as a multiple of ``embed_dim``. + norm_layer: Normalisation layer constructor. + attention_type: One of ``windowed-alternating``, ``alternating``, ``full``. + drop_path: Stochastic depth rate. + attn_drop: Dropout on attention weights. + proj_drop: Dropout on projection and feed-forward outputs. + max_channels: Channel capacity of the temporal position table. + max_timesteps: Timestep capacity used to size the temporal position table. + window_size_spatial: Channel-axis window size. + window_size_temporal: Time-axis window size. + dilation_cycle_spatial: Per-pair channel-axis dilations, cycled over blocks. + dilation_cycle_temporal: Per-pair time-axis dilations, cycled over blocks. + shift_cycle_spatial: Per-pair channel-axis window shifts, cycled over blocks. + shift_cycle_temporal: Per-pair time-axis window shifts, cycled over blocks. + use_axial_mode: Run all spatial blocks before all temporal blocks instead of + alternating them. + """ + + def __init__( + self, + patch_size: int = 200, + num_channels: int = 64, + embed_dim: int = 200, + depth: int = 12, + num_heads: int = 10, + mlp_ratio: float = 4.0, + norm_layer: nn.Module = nn.LayerNorm, + attention_type: str = "windowed-alternating", + drop_path: float = 0.0, + attn_drop: float = 0.1, + proj_drop: float = 0.1, + max_channels: int = 64, + max_timesteps: int = 6000, + window_size_spatial: int = 5, + window_size_temporal: int = 5, + dilation_cycle_spatial: Iterable[int] = (1, 2, 4), + dilation_cycle_temporal: Iterable[int] = (1, 2, 4), + shift_cycle_spatial: Iterable[int] = (-1, 1, -2, 2), + shift_cycle_temporal: Iterable[int] = (-1, 1, -2, 2), + use_axial_mode: bool = False, + ): + super().__init__() + self.patch_size = patch_size + self.num_channels = num_channels + self.embed_dim = embed_dim + self.depth = depth + self.num_heads = num_heads + self.max_channels = max_channels + self.max_timesteps = max_timesteps + self.max_patches = max_timesteps // patch_size + + if num_channels > max_channels: + raise ValueError(f"num_channels ({num_channels}) exceeds max_channels ({max_channels})") + + self.patch_embed = PatchEmbedding(patch_size=patch_size, embed_dim=embed_dim) + self.positional_embedding = PositionalEmbedding(self.max_patches, max_channels, embed_dim) + self.channel_embedding = ChannelEmbedding(embed_dim) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pad_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + + self.blocks = nn.ModuleList([ + TransformerBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=True, + norm_layer=norm_layer, + attention_type=attention_type, + block_idx=idx, + total_blocks=depth, + spatial_first=True, + drop_path=drop_path, + attn_drop=attn_drop, + proj_drop=proj_drop, + num_channels=num_channels, + window_size_spatial=window_size_spatial, + window_size_temporal=window_size_temporal, + dilation_cycle_spatial=dilation_cycle_spatial, + dilation_cycle_temporal=dilation_cycle_temporal, + shift_cycle_spatial=shift_cycle_spatial, + shift_cycle_temporal=shift_cycle_temporal, + use_axial_mode=use_axial_mode, + ) + for idx in range(depth) + ]) + self.norm = norm_layer(embed_dim) + + self.initialize_weights() + + def initialize_weights(self) -> None: + """Initialise special tokens, then every submodule.""" + trunc_normal_(self.pad_token, std=0.02, a=-0.02, b=0.02) + trunc_normal_(self.mask_token, std=0.02, a=-0.02, b=0.02) + trunc_normal_(self.positional_embedding.pos_embedding, std=0.02, a=-0.02, b=0.02) + self.apply(self._init_weights) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.xavier_uniform_(module.weight) + if module.bias is not None: + nn.init.constant_(module.bias, 0) + elif isinstance(module, nn.LayerNorm): + nn.init.constant_(module.bias, 0) + nn.init.constant_(module.weight, 1.0) + + def forward( + self, + x: torch.Tensor, + channel_positions: torch.Tensor, + directly_input_tokens: bool = False, + attn_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Encode a batch of patched EEG waveforms. + + Args: + x: Patched waveform of shape ``(batch, num_channels, num_patches, patch_size)``, + or token embeddings of shape ``(batch, num_channels, num_patches, embed_dim)`` + when ``directly_input_tokens`` is True. + channel_positions: Electrode coordinates of shape ``(batch, num_channels, 2, 3)``. + directly_input_tokens: Skip patch embedding and treat ``x`` as tokens. Used by + pre-training, which masks tokens between embedding and encoding. + attn_mask: Optional ``(batch, num_tokens)`` mask, 1 for real and 0 for padded tokens. + + Returns: + Contextualised token embeddings of shape ``(batch, num_tokens, embed_dim)``. + """ + batch, channels, patches = x.shape[0], x.shape[1], x.shape[2] + + if channels != self.num_channels: + raise ValueError( + f"Input has {channels} channels but the encoder was built for {self.num_channels}" + ) + if patches > self.max_patches: + raise ValueError( + f"Input has {patches} patches per channel, exceeding the positional " + f"embedding capacity of {self.max_patches}" + ) + + if not directly_input_tokens: + x = self.patch_embed(x) + + x = x.view(batch, channels, patches, self.embed_dim) + + pos_embed = self.positional_embedding()[:, :channels, :patches, :] + chan_embed = self.channel_embedding(channel_positions).unsqueeze(2) + + x = (x + pos_embed + chan_embed).reshape(batch, channels * patches, self.embed_dim) + + for block in self.blocks: + x = block(x=x, attn_mask=attn_mask) + + return self.norm(x) diff --git a/requirements.txt b/requirements.txt index de7de00..e16154e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,11 +4,13 @@ omegaconf hydra-core scipy h5py +lmdb python-dateutil lightning tqdm pandas mne +asrpy nvidia-nccl-cu11 peft psutil diff --git a/tasks/classification_task.py b/tasks/classification_task.py new file mode 100644 index 0000000..0c60c5d --- /dev/null +++ b/tasks/classification_task.py @@ -0,0 +1,398 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +from typing import Any, Dict, Optional + +import hydra +import pytorch_lightning as pl +import torch +import torch.nn as nn +from biofoundation.core.batch import BatchRequirements, as_signal_batch, require_batch_fields +from biofoundation.core.checkpoints import SafetensorsCheckpointMixin +from biofoundation.model_registry import get_model_spec +from torchmetrics.classification import ( + AUROC, + Accuracy, + AveragePrecision, + CohenKappa, + F1Score, + Precision, + Recall, +) + + +def split_checkpoint_state_dict(checkpoint: Dict[str, Any]) -> Dict[str, Dict[str, torch.Tensor]]: + """Separate a Lightning checkpoint into encoder and head state dicts. + + Args: + checkpoint: Object returned by ``torch.load``; either a Lightning checkpoint + with a ``state_dict`` entry or a bare state dict. + + Returns: + Mapping with keys ``model`` and ``model_head``, each a prefix-stripped state dict. + """ + if isinstance(checkpoint, dict) and "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + return { + "model": { + k[len("model."):]: v + for k, v in state_dict.items() + if k.startswith("model.") and not k.startswith("model_head.") + }, + "model_head": { + k[len("model_head."):]: v for k, v in state_dict.items() if k.startswith("model_head.") + }, + } + return {"model": dict(checkpoint), "model_head": {}} + + +def freeze_pretraining_only_parameters(encoder: nn.Module) -> list: + """Disable gradients for encoder parameters that no fine-tuning forward pass uses. + + The mask and pad tokens exist for masked pre-training. A fine-tuning step never + substitutes them, so they receive no gradient. Under DistributedDataParallel with + find_unused_parameters=False the reducer waits for a gradient from every parameter + it tracks, so leaving them trainable hangs the first training step of a multi-GPU + run with no error message. Freezing them removes them from DDP's set entirely, + which is cheaper than enabling find_unused_parameters and walking the graph on + every step. + + LUNA disables its mask token for classification for the same reason; see + models/LUNA.py. + + Returns: + Names of the parameters that were frozen. + """ + frozen = [] + for name in ("mask_token", "pad_token"): + parameter = getattr(encoder, name, None) + if parameter is not None and parameter.requires_grad: + parameter.requires_grad = False + frozen.append(name) + return frozen + + +class ClassificationTask(SafetensorsCheckpointMixin, pl.LightningModule): + """Fine-tuning task for EEG classification. + + Supports two input layouts: + + * **Window classification** (TUAB, CHB-MIT, Neonate, PhysioNet-MI, SHU-MI, STEW, + Mumtaz, MentalArithmetic, SEED-V): each sample is one window of shape + ``(num_channels, num_timesteps)`` with a single label. Use with + :class:`~models.model_heads.mlp_classification_head.MlpClassificationHead`. + * **Sequence classification** (ISRUC): each sample is a sequence of consecutive + epochs of shape ``(sequence_length, num_channels, num_timesteps)`` with one label + per epoch. The sequence axis is folded into the batch so the encoder processes + epochs independently, and per-epoch labels are flattened to match. Use with + :class:`~models.model_heads.sequence_classification_head.SequenceClassificationHead`, + which restores the sequence axis internally. + + Args: + hparams: Full experiment configuration. + freeze_backbone: Train only the head, the patch embedding and the embeddings. + layerwise_lr_decay: Per-block learning-rate decay; ``1.0`` disables it. Earlier + blocks receive ``lr * decay ** (depth - 1 - block_idx)``. + head_lr_multiplier: Multiplier applied to the head learning rate when the + optimiser config does not set ``head_lr`` explicitly. + """ + + LABEL_METRICS = ("acc", "balanced_acc", "f1_score", "precision", "cohen_kappa") + SCORE_METRICS = ("auroc", "aupr") + + def __init__( + self, + hparams, + freeze_backbone: bool = False, + layerwise_lr_decay: float = 1.0, + head_lr_multiplier: float = 1.0, + ): + super().__init__() + self.save_hyperparameters(hparams) + self.model = hydra.utils.instantiate(self.hparams.model) + self.model_head = hydra.utils.instantiate(self.hparams.model_head) + self.criterion = hydra.utils.instantiate(self.hparams.criterion) + + self.num_classes = int(self.hparams.model_head.num_classes) + if self.num_classes < 2: + raise ValueError(f"num_classes must be at least 2, got {self.num_classes}") + self.task_type = "binary" if self.num_classes == 2 else "multiclass" + + self.freeze_backbone = freeze_backbone + self.layerwise_lr_decay = layerwise_lr_decay + self.head_lr_multiplier = head_lr_multiplier + self.patch_size = int(self.hparams.model.patch_size) + self.softmax = nn.Softmax(dim=1) + self.strict_loading = False + + family = self.hparams.get("model_family", None) + self.batch_requirements = ( + get_model_spec(family).batch_requirements if family else BatchRequirements() + ) + + self.train_metrics = self._build_metrics() + self.val_metrics = self._build_metrics() + self.test_metrics = self._build_metrics() + + freeze_pretraining_only_parameters(self.model) + + if self.freeze_backbone: + self._apply_backbone_freeze() + + def _metrics(self, split: str) -> nn.ModuleDict: + """Return the metric set for ``train``, ``val`` or ``test``.""" + return getattr(self, f"{split}_metrics") + + def _build_metrics(self) -> nn.ModuleDict: + """Create one metric set for a single evaluation split.""" + return nn.ModuleDict({ + "acc": Accuracy(task=self.task_type, num_classes=self.num_classes), + "balanced_acc": Recall(task="multiclass", num_classes=self.num_classes, average="macro"), + "f1_score": F1Score(task="multiclass", num_classes=self.num_classes, average="weighted"), + "precision": Precision(task="multiclass", num_classes=self.num_classes, average="micro"), + "cohen_kappa": CohenKappa(task=self.task_type, num_classes=self.num_classes), + "auroc": AUROC(task=self.task_type, num_classes=self.num_classes, average="macro"), + "aupr": AveragePrecision(task=self.task_type, num_classes=self.num_classes, average="macro"), + }) + + def _apply_backbone_freeze(self) -> None: + """Freeze encoder blocks while leaving tokenisation and embeddings trainable.""" + trainable = ("patch_embed", "channel_embedding", "positional_embedding") + for name, param in self.model.named_parameters(): + param.requires_grad = any(prefix in name for prefix in trainable) + + def on_after_batch_transfer(self, batch: Dict[str, Any], dataloader_idx: int) -> Dict[str, Any]: + """Patch the waveforms and, for sequence data, fold the sequence into the batch.""" + x = batch["input"] + + if x.dim() == 3: + batch_size, channels, _ = x.shape + batch["input"] = x.reshape(batch_size, channels, -1, self.patch_size) + return batch + + if x.dim() == 4: + batch_size, sequence_length, channels, _ = x.shape + batch["input"] = x.reshape(batch_size * sequence_length, channels, -1, self.patch_size) + batch["label"] = batch["label"].reshape(batch_size * sequence_length) + batch["channel_coords"] = ( + batch["channel_coords"] + .unsqueeze(1) + .expand(-1, sequence_length, -1, -1, -1) + .reshape(batch_size * sequence_length, channels, 2, 3) + ) + return batch + + raise ValueError(f"Expected input with 3 or 4 dimensions, got {x.dim()}") + + def forward(self, x: torch.Tensor, channel_positions: torch.Tensor) -> torch.Tensor: + """Encode a batch and return classification logits.""" + encoded = self.model( + x, channel_positions=channel_positions, directly_input_tokens=False, attn_mask=None + ) + return self.model_head(encoded) + + def _shared_step(self, batch: Dict[str, Any], split: str) -> torch.Tensor: + """Compute the loss and update the metrics for one batch.""" + require_batch_fields(batch, self.batch_requirements) + logits = self(batch["input"], batch["channel_coords"]) + labels = batch["label"] + if labels.dim() == 2: + labels = labels.argmax(1) + batch["label"] = labels + + loss = self.criterion(logits, batch) + + predictions = torch.argmax(logits, dim=1) + probabilities = self.softmax(logits) + scores = probabilities[:, 1] if self.num_classes == 2 else probabilities + + metrics = self._metrics(split) + for name in self.LABEL_METRICS: + metrics[name](predictions, labels) + for name in self.SCORE_METRICS: + metrics[name](scores, labels) + + self.log( + f"{split}_loss", + loss, + on_step=True, + on_epoch=True, + prog_bar=True, + logger=True, + sync_dist=True, + batch_size=labels.shape[0], + ) + return loss + + def _log_epoch_metrics(self, split: str) -> None: + """Log and reset every metric for one split at epoch end.""" + for name, metric in self._metrics(split).items(): + self.log( + f"{split}_{name}", metric, prog_bar=True, logger=True, sync_dist=True, + on_step=False, on_epoch=True, + ) + + def training_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one training step.""" + if self.freeze_backbone: + self.model.eval() + return self._shared_step(as_signal_batch(batch), "train") + + def validation_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one validation step.""" + return self._shared_step(as_signal_batch(batch), "val") + + def test_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one test step.""" + return self._shared_step(as_signal_batch(batch), "test") + + def on_train_epoch_end(self) -> None: + """Log aggregated training metrics.""" + self._log_epoch_metrics("train") + + def on_validation_epoch_end(self) -> None: + """Log aggregated validation metrics.""" + self._log_epoch_metrics("val") + + def on_test_epoch_end(self) -> None: + """Log aggregated test metrics.""" + self._log_epoch_metrics("test") + + def configure_optimizers(self) -> Dict[str, Any]: + """Build parameter groups with layer-wise decay, then the optimiser and scheduler. + + Encoder blocks receive geometrically decayed learning rates so that layers + closer to the input change least. Biases, normalisation weights and embedding + tables are excluded from weight decay. The head forms its own group. + """ + base_lr = float(self.hparams.optimizer.lr) + base_weight_decay = float(getattr(self.hparams.optimizer, "weight_decay", 0.0)) + betas = tuple(getattr(self.hparams.optimizer, "betas", (0.9, 0.999))) + head_lr = float(getattr(self.hparams.optimizer, "head_lr", base_lr * self.head_lr_multiplier)) + head_weight_decay = float(getattr(self.hparams.optimizer, "head_weight_decay", base_weight_decay)) + depth = int(self.hparams.model.depth) + + param_groups = [] + for name, param in self.model.named_parameters(): + if not param.requires_grad: + continue + lr = base_lr + if self.layerwise_lr_decay != 1.0 and name.startswith("blocks."): + block_idx = int(name.split(".")[1]) + lr = base_lr * (self.layerwise_lr_decay ** (depth - 1 - block_idx)) + weight_decay = 0.0 if self._excluded_from_weight_decay(name, param) else base_weight_decay + param_groups.append({"params": [param], "lr": lr, "weight_decay": weight_decay}) + + head_params = [p for p in self.model_head.parameters() if p.requires_grad] + if head_params: + param_groups.append({"params": head_params, "lr": head_lr, "weight_decay": head_weight_decay}) + + optimizer_name = str(self.hparams.optimizer.optim).lower() + if optimizer_name == "adamw": + optimizer = torch.optim.AdamW(param_groups, lr=base_lr, weight_decay=base_weight_decay, betas=betas) + elif optimizer_name == "adam": + optimizer = torch.optim.Adam(param_groups, lr=base_lr, weight_decay=base_weight_decay, betas=betas) + elif optimizer_name == "sgd": + momentum = float(getattr(self.hparams.optimizer, "momentum", 0.9)) + optimizer = torch.optim.SGD( + param_groups, lr=base_lr, weight_decay=base_weight_decay, momentum=momentum + ) + else: + raise NotImplementedError(f"Unsupported optimizer: {self.hparams.optimizer.optim}") + + scheduler = hydra.utils.instantiate( + self.hparams.scheduler, + optimizer=optimizer, + total_training_opt_steps=self.trainer.estimated_stepping_batches, + ) + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1}, + } + + @staticmethod + def _excluded_from_weight_decay(name: str, param: torch.nn.Parameter) -> bool: + """Return True for parameters that should not be weight-decayed.""" + if name.endswith(".bias") or param.ndim == 1: + return True + lowered = name.lower() + return any( + key in lowered + for key in ("norm", "positional_embedding", "channel_embedding", "mask_token", "pad_token") + ) + + def lr_scheduler_step(self, scheduler, metric) -> None: + """Advance the timm-style scheduler once per optimiser step.""" + scheduler.step_update(num_updates=self.global_step) + + def load_from_checkpoint( + self, + checkpoint_path, + map_location=None, + hparams_file=None, + strict=None, + include_head: bool = False, + **kwargs, + ) -> "ClassificationTask": + """Load encoder weights, and optionally head weights, from a checkpoint. + + Tensors whose shapes do not match the current model are skipped rather than + forced, so an encoder pre-trained at a different channel count can seed + fine-tuning. The number of loaded and skipped tensors is printed so a silent + no-op load is visible in the logs. + """ + checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + state_dicts = split_checkpoint_state_dict(checkpoint) + + self._partial_load(self.model, state_dicts["model"], "model") + if include_head and state_dicts["model_head"]: + self._partial_load(self.model_head, state_dicts["model_head"], "model_head") + + if self.freeze_backbone: + self._apply_backbone_freeze() + return self + + @staticmethod + def _partial_load(module: nn.Module, incoming: Dict[str, torch.Tensor], label: str) -> None: + """Load only the tensors whose names and shapes match ``module``.""" + current = module.state_dict() + loaded, skipped, unexpected = [], [], [] + + for key, value in incoming.items(): + if key not in current: + unexpected.append(key) + elif value.shape != current[key].shape: + skipped.append(key) + else: + current[key] = value + loaded.append(key) + + module.load_state_dict(current, strict=False) + print( + f"[load:{label}] loaded={len(loaded)} shape_mismatch={len(skipped)} " + f"unexpected={len(unexpected)} total_target={len(current)}" + ) + if skipped: + print(f"[load:{label}] shape mismatch (first 10): {skipped[:10]}") + if unexpected: + print(f"[load:{label}] unexpected (first 10): {unexpected[:10]}") + if not loaded: + print(f"[load:{label}] WARNING: no tensors were loaded from this checkpoint") diff --git a/tasks/mae_pretraining.py b/tasks/mae_pretraining.py new file mode 100644 index 0000000..68fe6dd --- /dev/null +++ b/tasks/mae_pretraining.py @@ -0,0 +1,261 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +from typing import Any, Dict, Optional, Tuple + +import hydra +import pytorch_lightning as pl +import torch +import torch_optimizer as torch_optim + +from biofoundation.core.batch import BatchRequirements, as_signal_batch, require_batch_fields +from biofoundation.core.checkpoints import SafetensorsCheckpointMixin +from biofoundation.model_registry import get_model_spec +from models.modules.patching import patchify + + +def extract_encoder_state_dict(checkpoint: Dict[str, Any]) -> Dict[str, torch.Tensor]: + """Pull encoder weights out of a Lightning checkpoint. + + Args: + checkpoint: Checkpoint dictionary containing a ``state_dict`` entry. + + Returns: + State dict for the encoder, with the ``model.`` prefix removed and any + head parameters dropped. + """ + state_dict = checkpoint["state_dict"] + return { + key[len("model."):]: value + for key, value in state_dict.items() + if key.startswith("model.") and not key.startswith("model_head.") + } + + +class MaskedAutoencoderPretrainingTask(SafetensorsCheckpointMixin, pl.LightningModule): + """SimMIM-style masked reconstruction pre-training for the CEReBrO encoder. + + Waveforms are patched and embedded, a random subset of real (non-padded) tokens + is replaced by a learned mask token, and the encoder sees the full sequence of + visible and masked tokens. A linear decoder reconstructs every patch and the loss + is taken over the masked positions. Padded channels are replaced by a learned pad + token, are excluded from masking, and are masked out of attention. + + Args: + hparams: Full experiment configuration, used to instantiate the encoder, + decoder, criterion, optimiser and scheduler. + masking_ratio: Fraction of real tokens replaced by the mask token. + """ + + def __init__(self, hparams, masking_ratio: float = 0.5): + super().__init__() + self.save_hyperparameters(hparams) + self.model = hydra.utils.instantiate(self.hparams.model) + self.model_head = hydra.utils.instantiate(self.hparams.model_head) + self.criterion = hydra.utils.instantiate(self.hparams.criterion) + + family = self.hparams.get("model_family", None) + self.batch_requirements = ( + get_model_spec(family).batch_requirements if family else BatchRequirements() + ) + + self.masking_ratio = masking_ratio + self.patch_size = self.hparams.model.patch_size + self.num_channels = self.hparams.model.num_channels + self.embed_dim = self.hparams.model.embed_dim + # The mask and pad tokens are referenced through self.model rather than + # aliased onto this module. Assigning an nn.Parameter to a second module + # registers it a second time, so the checkpoint would carry both + # "model.mask_token" and "mask_token" backed by one storage. safetensors + # rejects shared storage, which would make the encoder undistributable in the + # format the Hugging Face releases use. + self.strict_loading = False + + def on_after_batch_transfer(self, batch: Dict[str, Any], dataloader_idx: int) -> Dict[str, Any]: + """Reshape raw waveforms into patches once the batch is on device.""" + x = batch["input"] + if x.dim() == 3: + batch_size, channels, _ = x.shape + batch["input"] = x.reshape(batch_size, channels, -1, self.patch_size) + return batch + + def _shared_step(self, batch: Dict[str, Any]) -> torch.Tensor: + """Run masking, encoding, decoding and loss for one batch.""" + require_batch_fields(batch, self.batch_requirements) + x = batch["input"] + batch_size, channels = x.shape[0], x.shape[1] + + tokens = self.model.patch_embed(x) + tokens, token_mask, attn_mask = self.prepare_tokens( + tokens, num_padded_channels=batch.get("num_padded_channels") + ) + + latent = self.model( + tokens, + channel_positions=batch["channel_coords"], + directly_input_tokens=True, + attn_mask=attn_mask, + ) + pred = self.model_head(latent) + + batch["token_mask"] = token_mask + batch["attn_mask"] = attn_mask + batch["target"] = patchify(x.reshape(batch_size, channels, -1), patch_size=self.patch_size) + + loss, _ = self.criterion(pred, batch) + return loss + + def training_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Compute and log the training reconstruction loss.""" + loss = self._shared_step(as_signal_batch(batch)) + self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) + return loss + + def validation_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Compute and log the validation reconstruction loss.""" + loss = self._shared_step(as_signal_batch(batch)) + self.log("val_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) + return loss + + def prepare_tokens( + self, tokens: torch.Tensor, num_padded_channels: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Insert pad tokens, then mask a random subset of the real tokens. + + Args: + tokens: Token embeddings of shape ``(batch, num_tokens, embed_dim)``. + num_padded_channels: Per-sample count of trailing padded channels. + + Returns: + Tuple of tokens reshaped to ``(batch, num_channels, num_patches, embed_dim)``, + a boolean ``(batch, num_tokens)`` mask marking masked tokens, and an + integer ``(batch, num_tokens)`` attention mask (``None`` when nothing is padded). + """ + batch_size, num_tokens, embed_dim = tokens.shape + channels = self.num_channels + patches = num_tokens // channels + + attn_mask = None + if num_padded_channels is not None: + num_real_channels = channels - num_padded_channels + channel_indices = torch.arange(channels, device=tokens.device).unsqueeze(0) + padded = channel_indices >= num_real_channels.unsqueeze(1) + padded = padded.repeat_interleave(patches, dim=1) + attn_mask = (~padded).int() + tokens = torch.where(padded.unsqueeze(-1), self.model.pad_token.to(tokens.dtype), tokens) + + tokens, token_mask = self.mask_tokens(tokens, attn_mask) + return tokens.reshape(batch_size, channels, patches, embed_dim), token_mask, attn_mask + + def mask_tokens( + self, tokens: torch.Tensor, attn_mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Replace a random subset of real tokens with the learned mask token. + + Tokens are ranked by uniform noise, with padded positions pushed to the end so + they are never selected. The first ``1 - masking_ratio`` fraction of each + sample's real tokens is kept and the remainder is masked in place, so the + sequence order the encoder sees is unchanged. + + Args: + tokens: Token embeddings of shape ``(batch, num_tokens, embed_dim)``. + attn_mask: Optional ``(batch, num_tokens)`` mask, 1 for real tokens. + + Returns: + Tuple of the masked tokens and a boolean ``(batch, num_tokens)`` mask + where True marks a masked token. + """ + batch_size, num_tokens, _ = tokens.shape + device = tokens.device + + noise = torch.rand(batch_size, num_tokens, device=device) + if attn_mask is not None: + noise = noise.masked_fill(attn_mask == 0, 2.0) + valid_length = attn_mask.sum(dim=1) + else: + valid_length = torch.full((batch_size,), num_tokens, device=device, dtype=torch.long) + + rank = torch.argsort(torch.argsort(noise, dim=1), dim=1) + num_keep = (valid_length * (1 - self.masking_ratio)).to(torch.long) + token_mask = (rank >= num_keep.unsqueeze(1)) & (rank < valid_length.unsqueeze(1)) + + masked = torch.where(token_mask.unsqueeze(-1), self.model.mask_token.to(tokens.dtype), tokens) + return masked, token_mask + + def configure_optimizers(self) -> Dict[str, Any]: + """Build the optimiser and the per-step learning-rate scheduler.""" + params = list(self.model.parameters()) + list(self.model_head.parameters()) + optimizer_name = str(self.hparams.optimizer.optim).lower() + lr = float(self.hparams.optimizer.lr) + weight_decay = float(getattr(self.hparams.optimizer, "weight_decay", 0.0)) + + if optimizer_name == "adamw": + betas = tuple(getattr(self.hparams.optimizer, "betas", (0.9, 0.999))) + optimizer = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay, betas=betas) + elif optimizer_name == "adam": + optimizer = torch.optim.Adam(params, lr=lr, weight_decay=weight_decay) + elif optimizer_name == "sgd": + momentum = float(getattr(self.hparams.optimizer, "momentum", 0.9)) + optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum, weight_decay=weight_decay) + elif optimizer_name == "lamb": + optimizer = torch_optim.Lamb(params, lr=lr) + else: + raise NotImplementedError(f"Unsupported optimizer: {self.hparams.optimizer.optim}") + + scheduler = hydra.utils.instantiate( + self.hparams.scheduler, + optimizer=optimizer, + total_training_opt_steps=self.trainer.estimated_stepping_batches, + ) + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1}, + } + + def lr_scheduler_step(self, scheduler, metric) -> None: + """Advance the timm-style scheduler once per optimiser step.""" + scheduler.step_update(num_updates=self.global_step) + + def load_from_checkpoint( + self, checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs + ) -> "MaskedAutoencoderPretrainingTask": + """Load encoder weights from a checkpoint, skipping the decoder. + + Shape-mismatched tensors are left at their initialised values so that an + encoder pre-trained at one channel count or window length can seed another. + """ + checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + incoming = extract_encoder_state_dict(checkpoint) + current = self.model.state_dict() + + loaded, skipped = [], [] + for key, value in incoming.items(): + if key in current and value.shape == current[key].shape: + current[key] = value + loaded.append(key) + else: + skipped.append(key) + + self.model.load_state_dict(current, strict=False) + print(f"[load] encoder tensors loaded: {len(loaded)}, skipped: {len(skipped)}") + if skipped: + print(f"[load] skipped keys (first 10): {skipped[:10]}") + return self diff --git a/tasks/regression_task.py b/tasks/regression_task.py new file mode 100644 index 0000000..341c8ce --- /dev/null +++ b/tasks/regression_task.py @@ -0,0 +1,305 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +from typing import Any, Dict + +import hydra +import pytorch_lightning as pl +import torch +import torch.nn as nn +from biofoundation.core.batch import BatchRequirements, as_signal_batch, require_batch_fields +from biofoundation.core.checkpoints import SafetensorsCheckpointMixin +from biofoundation.model_registry import get_model_spec +from torchmetrics import Metric +from torchmetrics.regression import MeanSquaredError, PearsonCorrCoef, R2Score + +from tasks.classification_task import ( + freeze_pretraining_only_parameters, + split_checkpoint_state_dict, +) + + +class NormalizedRootMeanSquaredError(Metric): + """Root mean squared error divided by the spread of the targets. + + Normalising by the target standard deviation (or mean) makes the error + comparable across datasets with different target scales. Statistics are + accumulated over the whole evaluation split rather than per batch. + + Args: + normalization: ``std`` to divide by the target standard deviation, ``mean`` to + divide by the absolute target mean. + """ + + def __init__(self, normalization: str = "std"): + super().__init__() + if normalization not in {"std", "mean"}: + raise ValueError(f"normalization must be 'std' or 'mean', got '{normalization}'") + self.normalization = normalization + self.add_state("sum_squared_error", default=torch.tensor(0.0), dist_reduce_fx="sum") + self.add_state("sum_target", default=torch.tensor(0.0), dist_reduce_fx="sum") + self.add_state("sum_squared_target", default=torch.tensor(0.0), dist_reduce_fx="sum") + self.add_state("num_observations", default=torch.tensor(0), dist_reduce_fx="sum") + + def update(self, preds: torch.Tensor, target: torch.Tensor) -> None: + """Accumulate error and target statistics for one batch.""" + preds = preds.float().flatten() + target = target.float().flatten() + self.sum_squared_error += torch.sum((preds - target) ** 2) + self.sum_target += torch.sum(target) + self.sum_squared_target += torch.sum(target ** 2) + self.num_observations += target.numel() + + def compute(self) -> torch.Tensor: + """Return the normalised RMSE over everything accumulated so far.""" + rmse = torch.sqrt(self.sum_squared_error / self.num_observations) + mean_target = self.sum_target / self.num_observations + + if self.normalization == "mean": + norm = torch.abs(mean_target) + else: + variance = (self.sum_squared_target / self.num_observations) - mean_target ** 2 + norm = torch.sqrt(torch.clamp(variance, min=0.0)) + + return rmse / torch.clamp(norm, min=1e-8) + + +class RegressionTask(SafetensorsCheckpointMixin, pl.LightningModule): + """Fine-tuning task for scalar EEG regression. + + Each sample is one window of shape ``(num_channels, num_timesteps)`` with a single + continuous target. Used for SEED-VIG vigilance (PERCLOS) estimation. Reported + metrics are RMSE, normalised RMSE, R² and Pearson correlation. + + Args: + hparams: Full experiment configuration. + freeze_backbone: Train only the head, the patch embedding and the embeddings. + layerwise_lr_decay: Per-block learning-rate decay; ``1.0`` disables it. + """ + + METRIC_NAMES = ("rmse", "nrmse", "r2", "pearson") + + def __init__(self, hparams, freeze_backbone: bool = False, layerwise_lr_decay: float = 1.0): + super().__init__() + self.save_hyperparameters(hparams) + self.model = hydra.utils.instantiate(self.hparams.model) + self.model_head = hydra.utils.instantiate(self.hparams.model_head) + self.criterion = hydra.utils.instantiate(self.hparams.criterion) + + family = self.hparams.get("model_family", None) + self.batch_requirements = ( + get_model_spec(family).batch_requirements if family else BatchRequirements() + ) + + self.freeze_backbone = freeze_backbone + self.layerwise_lr_decay = layerwise_lr_decay + self.patch_size = int(self.hparams.model.patch_size) + self.strict_loading = False + + self.train_metrics = self._build_metrics() + self.val_metrics = self._build_metrics() + self.test_metrics = self._build_metrics() + + freeze_pretraining_only_parameters(self.model) + + if self.freeze_backbone: + self._apply_backbone_freeze() + + def _metrics(self, split: str) -> nn.ModuleDict: + """Return the metric set for ``train``, ``val`` or ``test``.""" + return getattr(self, f"{split}_metrics") + + @staticmethod + def _build_metrics() -> nn.ModuleDict: + """Create one metric set for a single evaluation split.""" + return nn.ModuleDict({ + "rmse": MeanSquaredError(squared=False), + "nrmse": NormalizedRootMeanSquaredError(normalization="std"), + "r2": R2Score(), + "pearson": PearsonCorrCoef(num_outputs=1), + }) + + def _apply_backbone_freeze(self) -> None: + """Freeze encoder blocks while leaving tokenisation and embeddings trainable.""" + trainable = ("patch_embed", "channel_embedding", "positional_embedding") + for name, param in self.model.named_parameters(): + param.requires_grad = any(prefix in name for prefix in trainable) + + def on_after_batch_transfer(self, batch: Dict[str, Any], dataloader_idx: int) -> Dict[str, Any]: + """Reshape raw waveforms into patches once the batch is on device.""" + x = batch["input"] + if x.dim() != 3: + raise ValueError(f"Expected input with 3 dimensions, got {x.dim()}") + batch_size, channels, _ = x.shape + batch["input"] = x.reshape(batch_size, channels, -1, self.patch_size) + return batch + + def forward(self, x: torch.Tensor, channel_positions: torch.Tensor) -> torch.Tensor: + """Encode a batch and return scalar predictions.""" + encoded = self.model( + x, channel_positions=channel_positions, directly_input_tokens=False, attn_mask=None + ) + return self.model_head(encoded) + + def _shared_step(self, batch: Dict[str, Any], split: str) -> torch.Tensor: + """Compute the loss and update the metrics for one batch.""" + require_batch_fields(batch, self.batch_requirements) + predictions = self(batch["input"], batch["channel_coords"]) + targets = batch["label"].to(predictions.dtype).reshape(predictions.shape) + batch["label"] = targets + + loss = self.criterion(predictions, batch) + + for metric in self._metrics(split).values(): + metric(predictions, targets) + + self.log( + f"{split}_loss", + loss, + on_step=True, + on_epoch=True, + prog_bar=True, + logger=True, + sync_dist=True, + batch_size=targets.shape[0], + ) + return loss + + def _log_epoch_metrics(self, split: str) -> None: + """Log and reset every metric for one split at epoch end.""" + for name, metric in self._metrics(split).items(): + self.log( + f"{split}_{name}", metric, prog_bar=True, logger=True, sync_dist=True, + on_step=False, on_epoch=True, + ) + + def training_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one training step.""" + if self.freeze_backbone: + self.model.eval() + return self._shared_step(as_signal_batch(batch), "train") + + def validation_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one validation step.""" + return self._shared_step(as_signal_batch(batch), "val") + + def test_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one test step.""" + return self._shared_step(as_signal_batch(batch), "test") + + def on_train_epoch_end(self) -> None: + """Log aggregated training metrics.""" + self._log_epoch_metrics("train") + + def on_validation_epoch_end(self) -> None: + """Log aggregated validation metrics.""" + self._log_epoch_metrics("val") + + def on_test_epoch_end(self) -> None: + """Log aggregated test metrics.""" + self._log_epoch_metrics("test") + + def configure_optimizers(self) -> Dict[str, Any]: + """Build parameter groups with layer-wise decay, then the optimiser and scheduler.""" + base_lr = float(self.hparams.optimizer.lr) + base_weight_decay = float(getattr(self.hparams.optimizer, "weight_decay", 0.0)) + betas = tuple(getattr(self.hparams.optimizer, "betas", (0.9, 0.999))) + head_lr = float(getattr(self.hparams.optimizer, "head_lr", base_lr)) + head_weight_decay = float(getattr(self.hparams.optimizer, "head_weight_decay", base_weight_decay)) + depth = int(self.hparams.model.depth) + + param_groups = [] + for name, param in self.model.named_parameters(): + if not param.requires_grad: + continue + lr = base_lr + if self.layerwise_lr_decay != 1.0 and name.startswith("blocks."): + block_idx = int(name.split(".")[1]) + lr = base_lr * (self.layerwise_lr_decay ** (depth - 1 - block_idx)) + weight_decay = 0.0 if self._excluded_from_weight_decay(name, param) else base_weight_decay + param_groups.append({"params": [param], "lr": lr, "weight_decay": weight_decay}) + + head_params = [p for p in self.model_head.parameters() if p.requires_grad] + if head_params: + param_groups.append({"params": head_params, "lr": head_lr, "weight_decay": head_weight_decay}) + + optimizer_name = str(self.hparams.optimizer.optim).lower() + if optimizer_name == "adamw": + optimizer = torch.optim.AdamW(param_groups, lr=base_lr, weight_decay=base_weight_decay, betas=betas) + elif optimizer_name == "adam": + optimizer = torch.optim.Adam(param_groups, lr=base_lr, weight_decay=base_weight_decay, betas=betas) + elif optimizer_name == "sgd": + momentum = float(getattr(self.hparams.optimizer, "momentum", 0.9)) + optimizer = torch.optim.SGD( + param_groups, lr=base_lr, weight_decay=base_weight_decay, momentum=momentum + ) + else: + raise NotImplementedError(f"Unsupported optimizer: {self.hparams.optimizer.optim}") + + scheduler = hydra.utils.instantiate( + self.hparams.scheduler, + optimizer=optimizer, + total_training_opt_steps=self.trainer.estimated_stepping_batches, + ) + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step", "frequency": 1}, + } + + @staticmethod + def _excluded_from_weight_decay(name: str, param: torch.nn.Parameter) -> bool: + """Return True for parameters that should not be weight-decayed.""" + if name.endswith(".bias") or param.ndim == 1: + return True + lowered = name.lower() + return any( + key in lowered + for key in ("norm", "positional_embedding", "channel_embedding", "mask_token", "pad_token") + ) + + def lr_scheduler_step(self, scheduler, metric) -> None: + """Advance the timm-style scheduler once per optimiser step.""" + scheduler.step_update(num_updates=self.global_step) + + def load_from_checkpoint( + self, checkpoint_path, map_location=None, hparams_file=None, strict=None, **kwargs + ) -> "RegressionTask": + """Load encoder weights from a checkpoint, skipping the head.""" + checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False) + state_dicts = split_checkpoint_state_dict(checkpoint) + current = self.model.state_dict() + + loaded, skipped = [], [] + for key, value in state_dicts["model"].items(): + if key in current and value.shape == current[key].shape: + current[key] = value + loaded.append(key) + else: + skipped.append(key) + + self.model.load_state_dict(current, strict=False) + print(f"[load:model] loaded={len(loaded)} skipped={len(skipped)} total_target={len(current)}") + if not loaded: + print("[load:model] WARNING: no tensors were loaded from this checkpoint") + + if self.freeze_backbone: + self._apply_backbone_freeze() + return self diff --git a/tests/model_tests/test_attention.py b/tests/model_tests/test_attention.py new file mode 100644 index 0000000..26dff3e --- /dev/null +++ b/tests/model_tests/test_attention.py @@ -0,0 +1,131 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch + +from models.modules.attention import ( + AlternatingAttention, + WindowedAlternatingAttention, + build_window_indices, +) + + +def test_window_indices_are_centred_and_clamped(): + """A centred window of size 3 keeps interior neighbours and clamps at the edges.""" + indices = build_window_indices( + length=5, window_size=3, dilation=1, include_self=True, shift=0, device=torch.device("cpu") + ) + assert indices.shape == (5, 3) + assert indices[2].tolist() == [1, 2, 3] + assert indices[0].tolist() == [0, 0, 1] + assert indices[4].tolist() == [3, 4, 4] + + +def test_window_indices_apply_dilation_and_shift(): + """Dilation spaces the offsets out and shift translates the whole window.""" + dilated = build_window_indices( + length=9, window_size=3, dilation=2, include_self=True, shift=0, device=torch.device("cpu") + ) + assert dilated[4].tolist() == [2, 4, 6] + + shifted = build_window_indices( + length=9, window_size=3, dilation=1, include_self=True, shift=2, device=torch.device("cpu") + ) + assert shifted[4].tolist() == [5, 6, 7] + + +def test_window_size_is_clamped_to_axis_length(): + """A window larger than the axis degenerates to the whole axis instead of failing.""" + indices = build_window_indices( + length=3, window_size=11, dilation=1, include_self=True, shift=0, device=torch.device("cpu") + ) + assert indices.shape == (3, 3) + + +def test_full_window_matches_unwindowed_attention_on_unclamped_query(): + """With a window spanning every channel, the centre query matches full attention. + + Only the centre query is compared: at the edges the window is clamped, which + duplicates keys and legitimately changes the softmax. + """ + torch.manual_seed(0) + channels, patches, dim, heads = 7, 3, 16, 4 + + windowed = WindowedAlternatingAttention( + dim=dim, num_heads=heads, num_channels=channels, block_idx=0, + window_size_spatial=channels, dilation_spatial=1, include_self=True, + qkv_bias=True, qk_norm=False, + ).eval() + unwindowed = AlternatingAttention( + dim=dim, num_heads=heads, num_channels=channels, block_idx=0, qkv_bias=True, qk_norm=False, + ).eval() + unwindowed.load_state_dict(windowed.state_dict(), strict=False) + + x = torch.randn(2, channels * patches, dim) + with torch.no_grad(): + got = windowed(x).view(2, channels, patches, dim) + want = unwindowed(x).view(2, channels, patches, dim) + + centre = (channels - 1) // 2 + torch.testing.assert_close(got[:, centre], want[:, centre], atol=1e-5, rtol=1e-5) + + +def test_alternating_schedule_flips_axis_between_blocks(): + """Even blocks attend across channels and odd blocks across time.""" + even = WindowedAlternatingAttention(dim=8, num_heads=2, num_channels=4, block_idx=0) + odd = WindowedAlternatingAttention(dim=8, num_heads=2, num_channels=4, block_idx=1) + assert even.spatial_pass + assert not odd.spatial_pass + + +def test_axial_mode_splits_blocks_into_halves(): + """Axial mode runs every spatial block before every temporal block.""" + passes = [ + WindowedAlternatingAttention( + dim=8, num_heads=2, num_channels=4, block_idx=i, total_blocks=6, use_axial_mode=True + ).spatial_pass + for i in range(6) + ] + assert passes == [True, True, True, False, False, False] + + +def test_padded_queries_are_zeroed_and_padded_keys_are_ignored(): + """Padded tokens produce zero output and cannot influence real tokens.""" + torch.manual_seed(0) + channels, patches, dim = 6, 4, 16 + attention = WindowedAlternatingAttention( + dim=dim, num_heads=4, num_channels=channels, block_idx=0, window_size_spatial=channels + ).eval() + + x = torch.randn(1, channels * patches, dim) + mask = torch.ones(1, channels * patches, dtype=torch.int) + mask.view(1, channels, patches)[:, 4:] = 0 + + with torch.no_grad(): + out = attention(x, mask).view(1, channels, patches, dim) + + perturbed = x.clone().view(1, channels, patches, dim) + perturbed[:, 4:] = torch.randn_like(perturbed[:, 4:]) + out_perturbed = attention(perturbed.view(1, channels * patches, dim), mask) + out_perturbed = out_perturbed.view(1, channels, patches, dim) + + assert torch.all(out[:, 4:] == 0) + torch.testing.assert_close(out[:, :4], out_perturbed[:, :4], atol=1e-6, rtol=1e-6) diff --git a/tests/model_tests/test_pipeline.py b/tests/model_tests/test_pipeline.py new file mode 100644 index 0000000..82ed761 --- /dev/null +++ b/tests/model_tests/test_pipeline.py @@ -0,0 +1,204 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + +import torch +from omegaconf import OmegaConf + +from models.s_cerebro import SCerebroEncoder +from models.model_heads.mlp_classification_head import MlpClassificationHead +from models.model_heads.sequence_classification_head import SequenceClassificationHead +from models.modules.patching import patchify, unpatchify +from tasks.mae_pretraining import MaskedAutoencoderPretrainingTask + + +def build_encoder(**overrides): + """Construct a small encoder for shape and transfer tests.""" + kwargs = dict( + patch_size=200, num_channels=8, embed_dim=32, depth=4, num_heads=4, + max_channels=64, max_timesteps=6000, window_size_spatial=5, window_size_temporal=5, + ) + kwargs.update(overrides) + return SCerebroEncoder(**kwargs) + + +def test_patchify_roundtrip(): + """Patching then unpatching returns the original waveform.""" + signal = torch.randn(2, 5, 1200) + patches = patchify(signal, patch_size=200) + assert patches.shape == (2, 5 * 6, 200) + torch.testing.assert_close(unpatchify(patches, num_channels=5), signal) + + +def test_tokeniser_requires_the_supported_patch_size(): + """The tokeniser accepts 200-sample patches and rejects any other size.""" + encoder = build_encoder(patch_size=200) + x = torch.randn(1, 8, 4, 200) + coords = torch.randn(1, 8, 2, 3) + assert encoder(x, coords).shape == (1, 8 * 4, 32) + + for patch_size in (64, 128): + try: + build_encoder(patch_size=patch_size) + except ValueError as error: + assert "patch_size" in str(error) + else: + raise AssertionError(f"expected ValueError for patch_size={patch_size}") + + +def test_encoder_accepts_fewer_channels_and_patches_than_capacity(): + """A 6000-timestep position table serves a shorter window without reshaping.""" + encoder = build_encoder(num_channels=3, max_channels=64, max_timesteps=6000) + x = torch.randn(2, 3, 4, 200) + coords = torch.randn(2, 3, 2, 3) + assert encoder(x, coords).shape == (2, 12, 32) + + +def test_encoder_rejects_more_patches_than_capacity(): + """Exceeding the position table raises instead of silently truncating.""" + encoder = build_encoder(num_channels=2, max_timesteps=800) + x = torch.randn(1, 2, 8, 200) + coords = torch.randn(1, 2, 2, 3) + try: + encoder(x, coords) + except ValueError as error: + assert "positional embedding capacity" in str(error) + else: + raise AssertionError("expected ValueError for too many patches") + + +def test_pretrained_encoder_transfers_to_a_smaller_montage(): + """Every encoder tensor loads into a model built for fewer channels.""" + pretrained = build_encoder(num_channels=64) + finetuned = build_encoder(num_channels=8) + + source = pretrained.state_dict() + target = finetuned.state_dict() + mismatched = [k for k in source if k not in target or source[k].shape != target[k].shape] + assert mismatched == [] + + +def make_pretraining_task(masking_ratio=0.5, num_channels=8): + """Build a pre-training task from an in-memory config.""" + cfg = OmegaConf.create({ + "model": { + "_target_": "models.s_cerebro.SCerebroEncoder", + "patch_size": 200, "num_channels": num_channels, "embed_dim": 32, + "depth": 4, "num_heads": 4, "max_channels": 64, "max_timesteps": 6000, + }, + "model_head": { + "_target_": "models.model_heads.patch_reconstruction_head.PatchReconstructionHead", + "embed_dim": 32, "patch_size": 200, + }, + "criterion": { + "_target_": "criterion.masked_reconstruction_loss.MaskedReconstructionLoss", + "loss_type": "l2", "alpha": 0.1, + }, + }) + return MaskedAutoencoderPretrainingTask(cfg, masking_ratio=masking_ratio) + + +def test_token_mask_marks_exactly_the_replaced_tokens(): + """The reported mask lines up with the tokens that were actually replaced. + + The reconstruction loss is taken over positions flagged by ``token_mask``, so a + mask expressed in a different ordering than the tokens would score unmasked + patches and silently weaken the pre-training objective. + """ + torch.manual_seed(0) + task = make_pretraining_task() + tokens = torch.randn(4, 8 * 30, 32) + + masked, token_mask = task.mask_tokens(tokens.clone(), attn_mask=None) + + replaced = (masked == task.model.mask_token.detach()).all(dim=-1) + assert torch.equal(replaced, token_mask) + torch.testing.assert_close(masked[~token_mask], tokens[~token_mask]) + + +def test_masking_ratio_is_respected_and_padding_is_never_masked(): + """Masking hits the requested fraction of real tokens and no padded token.""" + torch.manual_seed(0) + channels, patches = 8, 30 + task = make_pretraining_task(masking_ratio=0.5, num_channels=channels) + tokens = torch.randn(2, channels * patches, 32) + + attn_mask = torch.ones(2, channels * patches, dtype=torch.int) + attn_mask.view(2, channels, patches)[:, 6:] = 0 + + _, token_mask = task.mask_tokens(tokens, attn_mask=attn_mask) + + real = attn_mask.sum(dim=1) + masked = token_mask.sum(dim=1) + assert torch.equal(masked, (real * 0.5).long()) + assert not token_mask[attn_mask == 0].any() + + +def test_sequence_head_maps_epochs_to_per_epoch_logits(): + """The ISRUC head restores the sequence axis and returns one logit row per epoch.""" + sequence_length, channels, patches, dim, classes = 20, 6, 30, 32, 5 + head = SequenceClassificationHead( + sequence_length=sequence_length, num_channels=channels, num_patches=patches, + embed_dim=dim, num_classes=classes, hidden_dim=64, dim_feedforward=128, + ) + encoded = torch.randn(2 * sequence_length, channels * patches, dim) + assert head(encoded).shape == (2 * sequence_length, classes) + + +def test_classification_head_pooling_modes_agree_on_output_shape(): + """Mean and flatten pooling both produce one logit row per window.""" + channels, patches, dim, classes = 4, 5, 32, 3 + encoded = torch.randn(2, channels * patches, dim) + for pooling_method in ("mean", "flatten"): + head = MlpClassificationHead( + embed_dim=dim, num_classes=classes, pooling_method=pooling_method, + num_channels=channels, num_patches=patches, + ) + assert head(encoded).shape == (2, classes) + + +def test_finetuning_leaves_no_trainable_parameter_without_a_gradient(): + """Every trainable parameter must receive a gradient from a fine-tuning step. + + DistributedDataParallel with find_unused_parameters=False waits for a gradient from + each parameter it tracks. A trainable parameter the forward pass never touches + therefore hangs the first training step of a multi-GPU run, with no error and no + output. The mask and pad tokens are pre-training-only and are frozen for exactly + this reason; this test fails if any other parameter joins them. + """ + import torch + + from models.model_heads.mlp_classification_head import MlpClassificationHead + from tasks.classification_task import freeze_pretraining_only_parameters + + encoder = build_encoder(num_channels=4, embed_dim=40, depth=2, num_heads=4) + freeze_pretraining_only_parameters(encoder) + head = MlpClassificationHead(embed_dim=40, num_classes=2, num_channels=4, num_patches=2) + + tokens = encoder(torch.randn(2, 4, 2, 200), channel_positions=torch.randn(2, 4, 2, 3)) + torch.nn.functional.cross_entropy(head(tokens), torch.tensor([0, 1])).backward() + + missing = [ + name + for module in (encoder, head) + for name, parameter in module.named_parameters() + if parameter.requires_grad and parameter.grad is None + ] + assert missing == [], f"trainable parameters with no gradient: {missing}" diff --git a/tests/test_batch.py b/tests/test_batch.py index ba995c8..634a0ad 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -57,6 +57,39 @@ def test_validates_model_specific_metadata(self): batch["sensor_type"] = object() self.assertIs(require_batch_fields(batch, requirements), batch) + def test_validates_paired_electrode_geometry_and_padding(self): + requirements = BatchRequirements(channel_coords=True, num_padded_channels=True) + batch = as_signal_batch({"input": object(), "channel_coords": object()}) + with self.assertRaisesRegex(ValueError, "num_padded_channels"): + require_batch_fields(batch, requirements) + + batch["num_padded_channels"] = object() + self.assertIs(require_batch_fields(batch, requirements), batch) + + def test_geometry_representations_are_independent(self): + """A model requiring one geometry field is not satisfied by the other.""" + + coords_only = as_signal_batch({"input": object(), "channel_coords": object()}) + with self.assertRaisesRegex(ValueError, "channel_locations"): + require_batch_fields(coords_only, BatchRequirements(channel_locations=True)) + + midpoints_only = as_signal_batch({"input": object(), "channel_locations": object()}) + with self.assertRaisesRegex(ValueError, "channel_coords"): + require_batch_fields(midpoints_only, BatchRequirements(channel_coords=True)) + + def test_defaults_leave_existing_requirements_unchanged(self): + """Fields added for newer families must not alter an existing spec's equality.""" + + self.assertEqual(BatchRequirements(), BatchRequirements()) + self.assertEqual( + BatchRequirements(channel_locations=True), + BatchRequirements(channel_locations=True), + ) + self.assertNotEqual( + BatchRequirements(channel_locations=True), + BatchRequirements(channel_coords=True), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_hydra_composition.py b/tests/test_hydra_composition.py index 1e754c1..c9d8ec3 100644 --- a/tests/test_hydra_composition.py +++ b/tests/test_hydra_composition.py @@ -17,6 +17,7 @@ #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* +import re import unittest from pathlib import Path @@ -33,6 +34,18 @@ ROOT = Path(__file__).resolve().parents[1] +# Experiments that do not compose today, with the reason. These are recorded rather +# than skipped silently so the breakage stays visible, and the test fails if one of +# them starts composing, which forces the entry to be removed rather than left to rot. +KNOWN_UNCOMPOSABLE_EXPERIMENTS = { + "FEMBA_quantized": ( + "pre-existing: its defaults list requires scheduler/constant_lr, which does " + "not exist in config/scheduler/. It also targets " + "ARES.tests.test_networks.test_24_femba_full_expland2, misspelling the module " + "test_24_femba_full_expand2." + ), +} + @unittest.skipIf(compose is None, "hydra-core is not installed") class HydraCompositionTest(unittest.TestCase): @@ -59,6 +72,102 @@ def test_every_registered_experiment_composes_and_resolves(self): self.assertIsInstance(resolved, dict) self.assertTrue(resolved["tag"]) + def test_every_experiment_file_composes_or_is_a_known_failure(self): + """Compose every experiment, including those no registry entry points at. + + The test above only reaches experiments named by a ModelSpec, which leaves + standalone experiment files unchecked. Composition failures that this catches + include defaults-list ordering mistakes, which produce a valid-looking YAML + file that Hydra rejects only at run time. + """ + + for path in sorted((ROOT / "config" / "experiment").glob("*.yaml")): + name = path.stem + with self.subTest(experiment=name): + try: + with initialize_config_dir(version_base="1.1", config_dir=str(ROOT / "config")): + config = compose(config_name="defaults", overrides=[f"+experiment={name}"]) + OmegaConf.to_container(config, resolve=True) + except Exception as error: # noqa: BLE001 - the failure itself is the assertion + if name in KNOWN_UNCOMPOSABLE_EXPERIMENTS: + self.skipTest(f"{name}: {KNOWN_UNCOMPOSABLE_EXPERIMENTS[name]}") + self.fail(f"{name} failed to compose: {type(error).__name__}: {error}") + else: + self.assertNotIn( + name, + KNOWN_UNCOMPOSABLE_EXPERIMENTS, + f"{name} now composes; remove it from KNOWN_UNCOMPOSABLE_EXPERIMENTS", + ) + + + def test_every_dataset_option_composes_with_a_consistent_head(self): + """Each config/dataset option must resolve and pair a head its task can drive. + + A dataset file owns its corpus path, sample layout, label kind, channel count, + prediction head, task and criterion together. Composing them separately is what + allowed a classification-only key such as num_classes to reach a regression + head, so this walks every option and checks the combination holds. + """ + dataset_dir = ROOT / "config" / "dataset" + options = sorted(path.stem for path in dataset_dir.glob("*.yaml")) + self.assertTrue(options, "no dataset options found") + + regression_heads = {"MlpRegressionHead"} + regression_tasks = {"RegressionTask"} + + for option in options: + with self.subTest(dataset=option): + with initialize_config_dir(version_base="1.1", config_dir=str(ROOT / "config")): + config = compose( + config_name="defaults", + overrides=["+experiment=SCEReBrO_finetune", f"dataset={option}"], + ) + resolved = OmegaConf.to_container(config, resolve=True) + + head = resolved["model_head"]["_target_"].rsplit(".", 1)[-1] + task = resolved["task"]["_target_"].rsplit(".", 1)[-1] + self.assertEqual( + head in regression_heads, + task in regression_tasks, + f"{option}: head {head} and task {task} disagree on regression", + ) + self.assertEqual( + resolved["label_mode"] == "regression", + task in regression_tasks, + f"{option}: label_mode and task disagree on regression", + ) + # A head only ever receives keys its constructor accepts. + if head in regression_heads: + self.assertNotIn("num_classes", resolved["model_head"], option) + self.assertNotIn("num_patches", resolved["model_head"], option) + else: + self.assertGreaterEqual(int(resolved["model_head"]["num_classes"]), 2, option) + + def test_finetuning_experiment_leaves_per_corpus_settings_to_the_dataset_group(self): + """The experiment must not restate anything config/dataset owns. + + Hydra applies a config's own values after its defaults list, so a key set in + both places resolves to the experiment's copy and the dataset file is silently + ignored. Keeping them out of the experiment makes one file the single owner. + """ + text = (ROOT / "config" / "experiment" / "SCEReBrO_finetune.yaml").read_text(encoding="utf-8") + # num_channels lives under model:, so leading whitespace is allowed for it. + owned = { + "dataset_root": r"^dataset_root:", + "dataset_kind": r"^dataset_kind:", + "label_mode": r"^label_mode:", + "model.num_channels": r"^\s*num_channels:", + } + restated = [key for key, pattern in owned.items() if re.search(pattern, text, re.MULTILINE)] + self.assertEqual(restated, [], f"restated in the experiment: {restated}") + + for group in ("model_head", "task", "criterion"): + self.assertNotRegex( + text, + rf"^\s*-\s*(override\s+)?/?{group}:", + f"{group} is selected by the experiment; leave it to config/dataset/", + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index fdb750a..0d1ebb4 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -30,17 +30,29 @@ class ModelRegistryTest(unittest.TestCase): def test_registry_contains_every_published_model(self): - self.assertEqual(set(MODEL_REGISTRY), {"femba", "luna", "tinymyo", "lumamba", "panluna"}) + self.assertEqual( + set(MODEL_REGISTRY), + {"femba", "luna", "tinymyo", "lumamba", "panluna", "s-cerebro"}, + ) + + def test_registry_keys_are_casefolded_display_names(self): + for key, spec in MODEL_REGISTRY.items(): + self.assertEqual(key, spec.display_name.casefold()) + + def _assert_target_resolves(self, target): + module_name, class_name = target.rsplit(".", 1) + module_path = ROOT / f"{module_name.replace('.', '/')}.py" + self.assertTrue(module_path.is_file(), module_path) + + tree = ast.parse(module_path.read_text(encoding="utf-8")) + classes = {node.name for node in tree.body if isinstance(node, ast.ClassDef)} + self.assertIn(class_name, classes, target) def test_model_targets_and_experiments_exist(self): for spec in MODEL_REGISTRY.values(): - module_name, class_name = spec.model_target.rsplit(".", 1) - module_path = ROOT / f"{module_name.replace('.', '/')}.py" - self.assertTrue(module_path.is_file(), module_path) - - tree = ast.parse(module_path.read_text(encoding="utf-8")) - classes = {node.name for node in tree.body if isinstance(node, ast.ClassDef)} - self.assertIn(class_name, classes, spec.model_target) + self._assert_target_resolves(spec.model_target) + for head_target in spec.head_targets: + self._assert_target_resolves(head_target) for experiment in (spec.pretrain_experiment, spec.finetune_experiment): path = ROOT / "config" / "experiment" / f"{experiment}.yaml" @@ -64,6 +76,10 @@ def test_batch_requirements_match_model_inputs(self): MODEL_REGISTRY["panluna"].batch_requirements, BatchRequirements(channel_locations=True, sensor_type=True), ) + self.assertEqual( + MODEL_REGISTRY["s-cerebro"].batch_requirements, + BatchRequirements(channel_coords=True), + ) def test_huggingface_and_paper_links_are_unique_and_documented(self): readme = (ROOT / "README.md").read_text(encoding="utf-8") @@ -78,16 +94,10 @@ def test_huggingface_and_paper_links_are_unique_and_documented(self): def test_citation_guide_covers_models_and_publication_venues(self): citations = (ROOT / "docs" / "CITATIONS.md").read_text(encoding="utf-8") - expected_venues = { - "FEMBA": "EMBC 2025", - "LUNA": "NeurIPS 2025", - "TinyMyo": "arXiv preprint", - "LuMamba": "EUSIPCO 2026", - "PanLUNA": "AICAS 2026", - } - for model, venue in expected_venues.items(): - self.assertIn(f"## {model}", citations) - self.assertIn(venue, citations) + for spec in MODEL_REGISTRY.values(): + self.assertTrue(spec.venue, f"{spec.display_name} has no venue in the registry") + self.assertIn(f"## {spec.display_name}", citations) + self.assertIn(spec.venue, citations) panluna_entry = citations.split("## PanLUNA", 1)[1] self.assertIn("Benini, Luca", panluna_entry) diff --git a/tests/test_repository_contracts.py b/tests/test_repository_contracts.py index b9f9a30..be9d672 100644 --- a/tests/test_repository_contracts.py +++ b/tests/test_repository_contracts.py @@ -25,18 +25,10 @@ ROOT = Path(__file__).resolve().parents[1] LOCAL_TARGET_PREFIXES = ("criterion.", "data_module.", "datasets.", "models.", "schedulers.", "tasks.") -TASK_FILES = ( - "finetune_regression_task_LuMamba.py", - "finetune_task.py", - "finetune_task_EMG.py", - "finetune_task_LUNA.py", - "finetune_task_PanLUNA.py", - "pretrain_task.py", - "pretrain_task_EMG.py", - "pretrain_task_LUNA.py", - "pretrain_task_LuMamba.py", - "pretrain_task_PanLUNA.py", -) + +# Derived rather than listed so that adding a task file cannot silently opt out of the +# shared batch-adapter contract below. +TASK_FILES = tuple(sorted(path.name for path in (ROOT / "tasks").glob("*.py"))) class RepositoryContractsTest(unittest.TestCase): @@ -114,6 +106,53 @@ def test_cli_validates_environment_before_hydra_starts(self): ] self.assertLess(call_names.index("require_environment"), call_names.index("run")) + def test_model_size_labels_match_their_config_filename(self): + """A model config that declares model_size must agree with its own filename. + + model_size names the output directory and is interpolated into pre-trained + checkpoint paths, so a config claiming a size it is not would silently load the + wrong weights. Declaring it in the model group rather than the experiment is + what keeps the two in step; this pins that the declaration is honest. + """ + pattern = re.compile(r"^model_size:\s*(\S+)", re.MULTILINE) + for config_path in (ROOT / "config" / "model").glob("*.yaml"): + match = pattern.search(config_path.read_text(encoding="utf-8")) + if match is None: + continue + expected = config_path.stem.rsplit("_", 1)[-1] + self.assertEqual(match.group(1), expected, config_path.name) + + def test_model_size_is_declared_in_exactly_one_place(self): + """An experiment must not restate a model_size its model group already declares. + + Families whose model configs carry model_size get it from whichever group is + selected, so the label always matches the encoder being built. Families that + never interpolate it may omit it entirely. Declaring it in both places is what + allows the label and the encoder to drift apart. + """ + declares = re.compile(r"^model_size:\s*(\S+)", re.MULTILINE) + selects_model = re.compile(r"^\s*-\s*override\s+/model:\s*(\S+)", re.MULTILINE) + + for experiment in sorted((ROOT / "config" / "experiment").glob("*.yaml")): + text = experiment.read_text(encoding="utf-8") + selected = selects_model.search(text) + if selected is None: + continue + model_config = ROOT / "config" / "model" / f"{selected.group(1)}.yaml" + if not model_config.is_file(): + continue + + group_declares = declares.search(model_config.read_text(encoding="utf-8")) is not None + experiment_declares = declares.search(text) is not None + # Not declaring it at all is fine: families that never interpolate + # model_size simply do not have the concept. + with self.subTest(experiment=experiment.name): + self.assertFalse( + group_declares and experiment_declares, + f"{experiment.name}: model_size declared in both the experiment and " + f"{model_config.name}; keep it in the model group only", + ) + def test_onboarding_docs_do_not_link_to_missing_local_paths(self): link_pattern = re.compile(r"\[[^\]]+\]\(([^)]+)\)") paths = (