From 89223a3a34d55ebf60a1deaccce0c22925dd0986 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 14:08:22 +0000 Subject: [PATCH 01/14] Extend biofoundation contracts for encoder/head model families Adds the shared contracts a separated encoder/head family needs, without altering anything the five bundled families observe. Every change is additive: new fields default to values that leave existing callers unchanged. biofoundation/core/batch.py SignalBatch gains channel_coords, num_padded_channels and num_padded_timesteps. channel_coords (batch, channels, 2, 3) keeps both electrodes of a channel separate; channel_locations (batch, channels, 3) stays as-is for the families that consume electrode midpoints. The two are documented as peers: a model requires one or the other, and neither is derived from the other. BatchRequirements gains matching flags, all defaulting to False, so require_batch_fields validates the new metadata only for families that ask for it. biofoundation/core/protocols.py New SignalEncoder and PredictionHead typing.Protocol definitions describing the token-embedding boundary. Conformance is structural, so bundled families remain valid unchanged while split families gain a checkable contract. Tensor annotations are guarded by TYPE_CHECKING to keep the package importable without PyTorch. biofoundation/core/checkpoints.py New SafetensorsCheckpointMixin maps the load_pretrained_checkpoint and load_safetensors_checkpoint entry points that run_train.py calls onto a task's own load_from_checkpoint, plus split_state_dict_by_prefix for separating encoder and head weights. biofoundation/model_registry.py ModelSpec gains venue, head_targets and size_variants, all defaulted. The five existing entries record their venue and, where applicable, their size variants. tests Venue coverage is now driven from the registry instead of a hard-coded map, and TASK_FILES is derived by globbing tasks/ so a new task file cannot opt out of the shared batch-adapter contract. New cases cover the paired-electrode and padding requirements, confirm the two geometry representations do not satisfy one another, and pin that the added defaults leave existing BatchRequirements equality unchanged. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- biofoundation/__init__.py | 18 ++++- biofoundation/core/__init__.py | 6 ++ biofoundation/core/batch.py | 35 ++++++++- biofoundation/core/checkpoints.py | 113 +++++++++++++++++++++++++++++ biofoundation/core/protocols.py | 81 +++++++++++++++++++++ biofoundation/model_registry.py | 18 ++++- tests/test_batch.py | 33 +++++++++ tests/test_model_registry.py | 33 ++++----- tests/test_repository_contracts.py | 16 +--- 9 files changed, 321 insertions(+), 32 deletions(-) create mode 100644 biofoundation/core/checkpoints.py create mode 100644 biofoundation/core/protocols.py 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..4e0c88f 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,7 @@ class ModelSpec: channel_locations=True, sensor_type=True, ), + venue="AICAS 2026", ), } ) 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_model_registry.py b/tests/test_model_registry.py index fdb750a..f7a0a26 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -32,15 +32,20 @@ class ModelRegistryTest(unittest.TestCase): def test_registry_contains_every_published_model(self): self.assertEqual(set(MODEL_REGISTRY), {"femba", "luna", "tinymyo", "lumamba", "panluna"}) + 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" @@ -78,16 +83,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..6a383e5 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): From 81e1efb5257cae581b590c35cadad93b5369599a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 14:09:43 +0000 Subject: [PATCH 02/14] Vendor the S-CEReBrO stack from TimeFM verbatim Imports the S-CEReBrO implementation with no edits of any kind, so that the adaptation that follows can be reviewed as a diff against known-good upstream code. Every file in this commit is byte-identical to its TimeFM counterpart and can be verified with cmp against the source tree. Imported: models/s_cerebro.py windowed alternating-attention encoder models/modules/{attention,patching,pos_embed}.py models/model_heads/*.py reconstruction, classification, regression and sequence heads criterion/{ce,focal,mse,masked_reconstruction}*.py tasks/{mae_pretraining,classification_task,regression_task}.py datasets/{lmdb_dataset,tueg_dataset}.py data_module/{pretraining,finetuning}_data_module.py make_datasets/*.py preprocessing for the pre-training corpus and 11 downstream datasets tests/model_tests/{test_attention,test_pipeline}.py Deliberately not imported: schedulers/cosine.py BioFoundation's CosineLRSchedulerWrapper already has a compatible constructor signature; the existing one is reused. callbacks/ run_train.py builds ModelCheckpoint from cfg.model_checkpoint, so NamedModelCheckpoint has no role here. config/, run_train.py replaced by BioFoundation-shaped equivalents in the adaptation commit. empty __init__.py copying these would turn existing BioFoundation namespace packages into regular packages. The upstream unit tests are placed under tests/model_tests/ rather than tests/. They are pytest-style and import torch, while the fast CI job installs neither. unittest discovery does not traverse a directory without __init__.py, so they stay invisible to `unittest discover -s tests` and runnable under pytest. Their contents are unmodified. This commit is intentionally inert: nothing references the new files yet. Two mechanical contract checks fail here by design and are resolved in the next commit, which is the first to edit any of this code: - Apache headers are absent from all 41 imported files. - The imported tasks do not yet call as_signal_batch in their step methods. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- criterion/ce_criterion.py | 34 ++ criterion/focal_criterion.py | 42 ++ criterion/masked_reconstruction_loss.py | 61 ++ criterion/mse_criterion.py | 25 + data_module/finetuning_data_module.py | 63 ++ data_module/pretraining_data_module.py | 77 +++ datasets/lmdb_dataset.py | 235 ++++++++ datasets/tueg_dataset.py | 125 ++++ make_datasets/common.py | 304 ++++++++++ make_datasets/electrode_positions.py | 158 +++++ make_datasets/make_bci_ner.py | 71 +++ make_datasets/make_boas.py | 62 ++ make_datasets/make_chb_mit.py | 136 +++++ make_datasets/make_gwd.py | 94 +++ make_datasets/make_isruc.py | 158 +++++ make_datasets/make_mental_arithmetic.py | 124 ++++ make_datasets/make_mumtaz.py | 126 ++++ make_datasets/make_neonate.py | 186 ++++++ make_datasets/make_physionet_mi.py | 114 ++++ make_datasets/make_seed.py | 70 +++ make_datasets/make_seed_fra_ger.py | 72 +++ make_datasets/make_seed_v.py | 128 +++++ make_datasets/make_seed_vig.py | 98 ++++ make_datasets/make_shu_mi.py | 78 +++ make_datasets/make_sleepedfx.py | 65 +++ make_datasets/make_stew.py | 146 +++++ make_datasets/make_tuab.py | 143 +++++ make_datasets/make_tueg.py | 181 ++++++ models/model_heads/mlp_classification_head.py | 97 ++++ models/model_heads/mlp_regression_head.py | 62 ++ .../model_heads/patch_reconstruction_head.py | 27 + .../sequence_classification_head.py | 109 ++++ models/modules/attention.py | 539 ++++++++++++++++++ models/modules/patching.py | 116 ++++ models/modules/pos_embed.py | 69 +++ models/s_cerebro.py | 187 ++++++ tasks/classification_task.py | 340 +++++++++++ tasks/mae_pretraining.py | 227 ++++++++ tasks/regression_task.py | 270 +++++++++ tests/model_tests/test_attention.py | 110 ++++ tests/model_tests/test_pipeline.py | 153 +++++ 41 files changed, 5482 insertions(+) create mode 100644 criterion/ce_criterion.py create mode 100644 criterion/focal_criterion.py create mode 100644 criterion/masked_reconstruction_loss.py create mode 100644 criterion/mse_criterion.py create mode 100644 data_module/finetuning_data_module.py create mode 100644 data_module/pretraining_data_module.py create mode 100644 datasets/lmdb_dataset.py create mode 100644 datasets/tueg_dataset.py create mode 100644 make_datasets/common.py create mode 100644 make_datasets/electrode_positions.py create mode 100644 make_datasets/make_bci_ner.py create mode 100644 make_datasets/make_boas.py create mode 100644 make_datasets/make_chb_mit.py create mode 100644 make_datasets/make_gwd.py create mode 100644 make_datasets/make_isruc.py create mode 100644 make_datasets/make_mental_arithmetic.py create mode 100644 make_datasets/make_mumtaz.py create mode 100644 make_datasets/make_neonate.py create mode 100644 make_datasets/make_physionet_mi.py create mode 100644 make_datasets/make_seed.py create mode 100644 make_datasets/make_seed_fra_ger.py create mode 100644 make_datasets/make_seed_v.py create mode 100644 make_datasets/make_seed_vig.py create mode 100644 make_datasets/make_shu_mi.py create mode 100644 make_datasets/make_sleepedfx.py create mode 100644 make_datasets/make_stew.py create mode 100644 make_datasets/make_tuab.py create mode 100644 make_datasets/make_tueg.py create mode 100644 models/model_heads/mlp_classification_head.py create mode 100644 models/model_heads/mlp_regression_head.py create mode 100644 models/model_heads/patch_reconstruction_head.py create mode 100644 models/model_heads/sequence_classification_head.py create mode 100644 models/modules/attention.py create mode 100644 models/modules/patching.py create mode 100644 models/modules/pos_embed.py create mode 100644 models/s_cerebro.py create mode 100644 tasks/classification_task.py create mode 100644 tasks/mae_pretraining.py create mode 100644 tasks/regression_task.py create mode 100644 tests/model_tests/test_attention.py create mode 100644 tests/model_tests/test_pipeline.py diff --git a/criterion/ce_criterion.py b/criterion/ce_criterion.py new file mode 100644 index 0000000..41f071f --- /dev/null +++ b/criterion/ce_criterion.py @@ -0,0 +1,34 @@ +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..376b76c --- /dev/null +++ b/criterion/focal_criterion.py @@ -0,0 +1,42 @@ +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..b924ff6 --- /dev/null +++ b/criterion/masked_reconstruction_loss.py @@ -0,0 +1,61 @@ +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..6e7cb06 --- /dev/null +++ b/criterion/mse_criterion.py @@ -0,0 +1,25 @@ +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..da7711b --- /dev/null +++ b/data_module/finetuning_data_module.py @@ -0,0 +1,63 @@ +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..cc451c5 --- /dev/null +++ b/data_module/pretraining_data_module.py @@ -0,0 +1,77 @@ +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..cc30c7a --- /dev/null +++ b/datasets/lmdb_dataset.py @@ -0,0 +1,235 @@ +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 = lmdb.open( + self.lmdb_path, readonly=True, lock=False, readahead=True, max_readers=64, map_async=True + ) + 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 _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] + + with self.env.begin(buffers=True) as txn: + return [bytes(key) for key in txn.cursor().iternext(keys=True, values=False)] + + 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] + + 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..321fd06 --- /dev/null +++ b/datasets/tueg_dataset.py @@ -0,0 +1,125 @@ +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/make_datasets/common.py b/make_datasets/common.py new file mode 100644 index 0000000..5b61924 --- /dev/null +++ b/make_datasets/common.py @@ -0,0 +1,304 @@ +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..09e2bb6 --- /dev/null +++ b/make_datasets/electrode_positions.py @@ -0,0 +1,158 @@ +"""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..e317d36 --- /dev/null +++ b/make_datasets/make_bci_ner.py @@ -0,0 +1,71 @@ +"""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..19722d1 --- /dev/null +++ b/make_datasets/make_boas.py @@ -0,0 +1,62 @@ +"""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..ca47d5c --- /dev/null +++ b/make_datasets/make_chb_mit.py @@ -0,0 +1,136 @@ +"""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..a44b957 --- /dev/null +++ b/make_datasets/make_gwd.py @@ -0,0 +1,94 @@ +"""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..99daeee --- /dev/null +++ b/make_datasets/make_isruc.py @@ -0,0 +1,158 @@ +"""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..2096e81 --- /dev/null +++ b/make_datasets/make_mental_arithmetic.py @@ -0,0 +1,124 @@ +"""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..5d31bb2 --- /dev/null +++ b/make_datasets/make_mumtaz.py @@ -0,0 +1,126 @@ +"""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..d0ccff7 --- /dev/null +++ b/make_datasets/make_neonate.py @@ -0,0 +1,186 @@ +"""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..c32e763 --- /dev/null +++ b/make_datasets/make_physionet_mi.py @@ -0,0 +1,114 @@ +"""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..20879ac --- /dev/null +++ b/make_datasets/make_seed.py @@ -0,0 +1,70 @@ +"""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..09f38ec --- /dev/null +++ b/make_datasets/make_seed_fra_ger.py @@ -0,0 +1,72 @@ +"""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..220d22e --- /dev/null +++ b/make_datasets/make_seed_v.py @@ -0,0 +1,128 @@ +"""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..775c08a --- /dev/null +++ b/make_datasets/make_seed_vig.py @@ -0,0 +1,98 @@ +"""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..e4bfd71 --- /dev/null +++ b/make_datasets/make_shu_mi.py @@ -0,0 +1,78 @@ +"""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..7163337 --- /dev/null +++ b/make_datasets/make_sleepedfx.py @@ -0,0 +1,65 @@ +"""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..509d173 --- /dev/null +++ b/make_datasets/make_stew.py @@ -0,0 +1,146 @@ +"""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..74aa6ee --- /dev/null +++ b/make_datasets/make_tuab.py @@ -0,0 +1,143 @@ +"""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..7d5eff4 --- /dev/null +++ b/make_datasets/make_tueg.py @@ -0,0 +1,181 @@ +"""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/model_heads/mlp_classification_head.py b/models/model_heads/mlp_classification_head.py new file mode 100644 index 0000000..dc9d9d1 --- /dev/null +++ b/models/model_heads/mlp_classification_head.py @@ -0,0 +1,97 @@ +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..aad7a56 --- /dev/null +++ b/models/model_heads/mlp_regression_head.py @@ -0,0 +1,62 @@ +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..591ab39 --- /dev/null +++ b/models/model_heads/patch_reconstruction_head.py @@ -0,0 +1,27 @@ +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..e3270cc --- /dev/null +++ b/models/model_heads/sequence_classification_head.py @@ -0,0 +1,109 @@ +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..29cd600 --- /dev/null +++ b/models/modules/attention.py @@ -0,0 +1,539 @@ +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..3c783ab --- /dev/null +++ b/models/modules/patching.py @@ -0,0 +1,116 @@ +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..09d233a --- /dev/null +++ b/models/modules/pos_embed.py @@ -0,0 +1,69 @@ +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..1d26bc0 --- /dev/null +++ b/models/s_cerebro.py @@ -0,0 +1,187 @@ +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/tasks/classification_task.py b/tasks/classification_task.py new file mode 100644 index 0000000..b47e22e --- /dev/null +++ b/tasks/classification_task.py @@ -0,0 +1,340 @@ +from typing import Any, Dict, Optional + +import hydra +import pytorch_lightning as pl +import torch +import torch.nn as nn +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": {}} + + +class ClassificationTask(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 + + self.train_metrics = self._build_metrics() + self.val_metrics = self._build_metrics() + self.test_metrics = self._build_metrics() + + 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.""" + 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(batch, "train") + + def validation_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one validation step.""" + return self._shared_step(batch, "val") + + def test_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one test step.""" + return self._shared_step(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..818dd72 --- /dev/null +++ b/tasks/mae_pretraining.py @@ -0,0 +1,227 @@ +from typing import Any, Dict, Optional, Tuple + +import hydra +import pytorch_lightning as pl +import torch +import torch_optimizer as torch_optim + +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(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) + + 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 + self.mask_token = self.model.mask_token + self.pad_token = self.model.pad_token + 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.""" + 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(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(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.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.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..f7e4a8f --- /dev/null +++ b/tasks/regression_task.py @@ -0,0 +1,270 @@ +from typing import Any, Dict + +import hydra +import pytorch_lightning as pl +import torch +import torch.nn as nn +from torchmetrics import Metric +from torchmetrics.regression import MeanSquaredError, PearsonCorrCoef, R2Score + +from tasks.classification_task import 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(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) + + 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() + + 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.""" + 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(batch, "train") + + def validation_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one validation step.""" + return self._shared_step(batch, "val") + + def test_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: + """Run one test step.""" + return self._shared_step(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..fde1314 --- /dev/null +++ b/tests/model_tests/test_attention.py @@ -0,0 +1,110 @@ +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..edd147d --- /dev/null +++ b/tests/model_tests/test_pipeline.py @@ -0,0 +1,153 @@ +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.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) From e50947bd92cb6c8de320ad9a7e2d6c88e3c4491f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 14:14:29 +0000 Subject: [PATCH 03/14] Adapt the vendored S-CEReBrO stack to BioFoundation conventions First commit to modify any imported file. Changes are confined to the repository's own conventions; no model, head, criterion or dataset logic is altered. Licensing Apache 2.0 headers added to all 41 imported files, each recording provenance from the S-CEReBrO reference implementation. New models/model_heads/__init__.py carries the header, matching models/modules/. Batch contract The three imported tasks now normalise input through as_signal_batch in every training, validation and test step, satisfying the shared adapter contract. Each task also validates its family's BatchRequirements once per step via require_batch_fields. The requirements are looked up from MODEL_REGISTRY using a model_family key set by the experiment, so the registry becomes load-bearing at runtime rather than metadata that only the tests read. Tasks fall back to empty requirements when no family is named, so they stay usable by future families. Checkpoints The tasks mix in SafetensorsCheckpointMixin, exposing the load_pretrained_checkpoint and load_safetensors_checkpoint entry points that run_train.py calls while keeping their own load_from_checkpoint as the single implementation. Its partial loader skips shape-mismatched tensors, so an encoder pre-trained at 64 channels can seed a 22-channel fine-tune. Hydra configuration config/model/SCEReBrO_{tiny,small,base}.yaml size variants, LUNA-style config/model_head/*.yaml new group: reconstruction, classification, regression, sequence config/task/{pretrain,finetune,finetune_regression}_task_SCEReBrO.yaml config/criterion/{masked_reconstruction_loss,ce,focal,mse}_criterion.yaml config/data_module/{pretrain,finetune}_data_module_SCEReBrO.yaml config/experiment/SCEReBrO_{pretrain,finetune}.yaml All follow the house convention: @package directive on line 1, licence header after it, contents under a wrapper key. Paths use ${env:DATA_PATH} and ${env:CHECKPOINT_DIR} rather than the upstream data_root defaults. model_head is deliberately absent from config/defaults.yaml and is instead added by the S-CEReBrO experiments through their own defaults list, so the composition graph of the ten existing experiments is untouched. Experiments are flat, matching the other five families. The fine-tuning corpus is a dataset_root override rather than a file per dataset, so all eleven prepared datasets are reachable from one experiment. The data module receives seed: ${seed} through config because run_train.py instantiates it without arguments, keeping the train/validation split identical across ranks. Dependencies lmdb and asrpy added to requirements.txt, the only imports in the vendored code not already satisfied. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- config/criterion/ce_criterion.yaml | 24 ++++ config/criterion/focal_criterion.yaml | 24 ++++ .../criterion/masked_reconstruction_loss.yaml | 24 ++++ config/criterion/mse_criterion.yaml | 22 ++++ .../finetune_data_module_SCEReBrO.yaml | 57 +++++++++ .../pretrain_data_module_SCEReBrO.yaml | 115 +++++++++++++++++ config/experiment/SCEReBrO_finetune.yaml | 119 ++++++++++++++++++ config/experiment/SCEReBrO_pretrain.yaml | 100 +++++++++++++++ config/model/SCEReBrO_base.yaml | 52 ++++++++ config/model/SCEReBrO_small.yaml | 52 ++++++++ config/model/SCEReBrO_tiny.yaml | 52 ++++++++ .../model_head/mlp_classification_head.yaml | 29 +++++ config/model_head/mlp_regression_head.yaml | 26 ++++ .../model_head/patch_reconstruction_head.yaml | 24 ++++ .../sequence_classification_head.yaml | 33 +++++ .../finetune_regression_task_SCEReBrO.yaml | 24 ++++ config/task/finetune_task_SCEReBrO.yaml | 25 ++++ config/task/pretrain_task_SCEReBrO.yaml | 23 ++++ criterion/ce_criterion.py | 22 ++++ criterion/focal_criterion.py | 22 ++++ criterion/masked_reconstruction_loss.py | 22 ++++ criterion/mse_criterion.py | 22 ++++ data_module/finetuning_data_module.py | 22 ++++ data_module/pretraining_data_module.py | 22 ++++ datasets/lmdb_dataset.py | 22 ++++ datasets/tueg_dataset.py | 22 ++++ make_datasets/common.py | 22 ++++ make_datasets/electrode_positions.py | 22 ++++ make_datasets/make_bci_ner.py | 22 ++++ make_datasets/make_boas.py | 22 ++++ make_datasets/make_chb_mit.py | 22 ++++ make_datasets/make_gwd.py | 22 ++++ make_datasets/make_isruc.py | 22 ++++ make_datasets/make_mental_arithmetic.py | 22 ++++ make_datasets/make_mumtaz.py | 22 ++++ make_datasets/make_neonate.py | 22 ++++ make_datasets/make_physionet_mi.py | 22 ++++ make_datasets/make_seed.py | 22 ++++ make_datasets/make_seed_fra_ger.py | 22 ++++ make_datasets/make_seed_v.py | 22 ++++ make_datasets/make_seed_vig.py | 22 ++++ make_datasets/make_shu_mi.py | 22 ++++ make_datasets/make_sleepedfx.py | 22 ++++ make_datasets/make_stew.py | 22 ++++ make_datasets/make_tuab.py | 22 ++++ make_datasets/make_tueg.py | 22 ++++ models/model_heads/__init__.py | 21 ++++ models/model_heads/mlp_classification_head.py | 22 ++++ models/model_heads/mlp_regression_head.py | 22 ++++ .../model_heads/patch_reconstruction_head.py | 22 ++++ .../sequence_classification_head.py | 22 ++++ models/modules/attention.py | 22 ++++ models/modules/patching.py | 22 ++++ models/modules/pos_embed.py | 22 ++++ models/s_cerebro.py | 22 ++++ requirements.txt | 2 + tasks/classification_task.py | 39 +++++- tasks/mae_pretraining.py | 37 +++++- tasks/regression_task.py | 39 +++++- tests/model_tests/test_attention.py | 22 ++++ tests/model_tests/test_pipeline.py | 22 ++++ 61 files changed, 1788 insertions(+), 11 deletions(-) create mode 100644 config/criterion/ce_criterion.yaml create mode 100644 config/criterion/focal_criterion.yaml create mode 100644 config/criterion/masked_reconstruction_loss.yaml create mode 100644 config/criterion/mse_criterion.yaml create mode 100644 config/data_module/finetune_data_module_SCEReBrO.yaml create mode 100644 config/data_module/pretrain_data_module_SCEReBrO.yaml create mode 100644 config/experiment/SCEReBrO_finetune.yaml create mode 100644 config/experiment/SCEReBrO_pretrain.yaml create mode 100644 config/model/SCEReBrO_base.yaml create mode 100644 config/model/SCEReBrO_small.yaml create mode 100644 config/model/SCEReBrO_tiny.yaml create mode 100644 config/model_head/mlp_classification_head.yaml create mode 100644 config/model_head/mlp_regression_head.yaml create mode 100644 config/model_head/patch_reconstruction_head.yaml create mode 100644 config/model_head/sequence_classification_head.yaml create mode 100644 config/task/finetune_regression_task_SCEReBrO.yaml create mode 100644 config/task/finetune_task_SCEReBrO.yaml create mode 100644 config/task/pretrain_task_SCEReBrO.yaml create mode 100644 models/model_heads/__init__.py diff --git a/config/criterion/ce_criterion.yaml b/config/criterion/ce_criterion.yaml new file mode 100644 index 0000000..c89eb81 --- /dev/null +++ b/config/criterion/ce_criterion.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: Glenn Anta Bucagu * +#* 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..c5d7cf4 --- /dev/null +++ b/config/criterion/focal_criterion.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: Glenn Anta Bucagu * +#* 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..9324c49 --- /dev/null +++ b/config/criterion/masked_reconstruction_loss.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: Glenn Anta Bucagu * +#* 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..e288140 --- /dev/null +++ b/config/criterion/mse_criterion.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: Glenn Anta Bucagu * +#* 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..5f8d6de --- /dev/null +++ b/config/data_module/finetune_data_module_SCEReBrO.yaml @@ -0,0 +1,57 @@ +# @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: Glenn Anta Bucagu * +#* 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..5756a98 --- /dev/null +++ b/config/data_module/pretrain_data_module_SCEReBrO.yaml @@ -0,0 +1,115 @@ +# @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: Glenn Anta Bucagu * +#* 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/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml new file mode 100644 index 0000000..d9ae46c --- /dev/null +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -0,0 +1,119 @@ +# @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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +tag: SCEReBrO_finetune +model_size: tiny + +# Names the registry entry whose BatchRequirements the task enforces on every batch. +model_family: scerebro + +# Fine-tuning dataset. Defaults to TUAB; switch corpus with a single override, e.g. +# python -u run_train.py +experiment=SCEReBrO_finetune \ +# dataset_root=${env:DATA_PATH}/finetuning/CHB-MIT model.num_channels=16 +# docs/model/SCEReBrO.md lists the channel count, window length and class count for +# each prepared dataset. +dataset_root: ${env:DATA_PATH}/finetuning/TUAB +dataset_kind: window # 'window', or 'sequence' for ISRUC sleep staging +label_mode: classification + +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: False + +pretrained_checkpoint_path: null +pretrained_safetensors_path: null + +defaults: + - override /data_module: finetune_data_module_SCEReBrO + - override /model: SCEReBrO_tiny # or SCEReBrO_small, SCEReBrO_base + - override /task: finetune_task_SCEReBrO + - override /criterion: ce_criterion + - override /scheduler: cosine + # For SEED-VIG vigilance regression, override these three together: + # /model_head: mlp_regression_head, /task: finetune_regression_task_SCEReBrO, + # /criterion: mse_criterion, and set label_mode=regression. + - /model_head: mlp_classification_head + +model: + num_channels: 22 # TUAB montage + +model_head: + num_classes: 2 + num_patches: 10 # 10 s window at one patch per second + +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..336f942 --- /dev/null +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -0,0 +1,100 @@ +# @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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +tag: SCEReBrO_pretrain +model_size: tiny + +# Names the registry entry whose BatchRequirements the task enforces on every batch. +model_family: scerebro + +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 + +defaults: + - 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_head is not part of the global defaults, so it is added here rather than + # overridden. The five bundled families compose exactly as before. + - /model_head: patch_reconstruction_head + +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: 50 + 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: 3 + 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..f9e0394 --- /dev/null +++ b/config/model/SCEReBrO_base.yaml @@ -0,0 +1,52 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +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..4a043b5 --- /dev/null +++ b/config/model/SCEReBrO_small.yaml @@ -0,0 +1,52 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +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..49f986c --- /dev/null +++ b/config/model/SCEReBrO_tiny.yaml @@ -0,0 +1,52 @@ +# @package _global_ +#*----------------------------------------------------------------------------* +#* Copyright (C) 2026 ETH Zurich, Switzerland * +#* SPDX-License-Identifier: Apache-2.0 * +#* * +#* Licensed under the Apache License, Version 2.0 (the "License"); * +#* you may not use this file except in compliance with the License. * +#* You may obtain a copy of the License at * +#* * +#* http://www.apache.org/licenses/LICENSE-2.0 * +#* * +#* Unless required by applicable law or agreed to in writing, software * +#* distributed under the License is distributed on an "AS IS" BASIS, * +#* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * +#* See the License for the specific language governing permissions and * +#* limitations under the License. * +#* * +#* Author: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#*----------------------------------------------------------------------------* +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..fadf8b7 --- /dev/null +++ b/config/model_head/mlp_classification_head.yaml @@ -0,0 +1,29 @@ +# @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: Glenn Anta Bucagu * +#* 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..63f3158 --- /dev/null +++ b/config/model_head/mlp_regression_head.yaml @@ -0,0 +1,26 @@ +# @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: Glenn Anta Bucagu * +#* 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..c07fd43 --- /dev/null +++ b/config/model_head/patch_reconstruction_head.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: Glenn Anta Bucagu * +#* 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..d793ca4 --- /dev/null +++ b/config/model_head/sequence_classification_head.yaml @@ -0,0 +1,33 @@ +# @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: Glenn Anta Bucagu * +#* 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..567920a --- /dev/null +++ b/config/task/finetune_regression_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: Glenn Anta Bucagu * +#* 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..fbb7acc --- /dev/null +++ b/config/task/finetune_task_SCEReBrO.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: Glenn Anta Bucagu * +#* 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..ad69397 --- /dev/null +++ b/config/task/pretrain_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: Glenn Anta Bucagu * +#* 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 index 41f071f..eed65a2 100644 --- a/criterion/ce_criterion.py +++ b/criterion/ce_criterion.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict, Optional, Sequence import torch diff --git a/criterion/focal_criterion.py b/criterion/focal_criterion.py index 376b76c..ee2e8c2 100644 --- a/criterion/focal_criterion.py +++ b/criterion/focal_criterion.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict import torch diff --git a/criterion/masked_reconstruction_loss.py b/criterion/masked_reconstruction_loss.py index b924ff6..69f82df 100644 --- a/criterion/masked_reconstruction_loss.py +++ b/criterion/masked_reconstruction_loss.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict, Tuple import torch diff --git a/criterion/mse_criterion.py b/criterion/mse_criterion.py index 6e7cb06..2aea9fa 100644 --- a/criterion/mse_criterion.py +++ b/criterion/mse_criterion.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict import torch diff --git a/data_module/finetuning_data_module.py b/data_module/finetuning_data_module.py index da7711b..e2de3b2 100644 --- a/data_module/finetuning_data_module.py +++ b/data_module/finetuning_data_module.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Optional import pytorch_lightning as pl diff --git a/data_module/pretraining_data_module.py b/data_module/pretraining_data_module.py index cc451c5..40e6ec7 100644 --- a/data_module/pretraining_data_module.py +++ b/data_module/pretraining_data_module.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict, Optional import pytorch_lightning as pl diff --git a/datasets/lmdb_dataset.py b/datasets/lmdb_dataset.py index cc30c7a..d62583f 100644 --- a/datasets/lmdb_dataset.py +++ b/datasets/lmdb_dataset.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import glob import json import os diff --git a/datasets/tueg_dataset.py b/datasets/tueg_dataset.py index 321fd06..146070c 100644 --- a/datasets/tueg_dataset.py +++ b/datasets/tueg_dataset.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import os from collections import OrderedDict from typing import Any, Dict, List, Optional diff --git a/make_datasets/common.py b/make_datasets/common.py index 5b61924..da3f8c3 100644 --- a/make_datasets/common.py +++ b/make_datasets/common.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import argparse import os import pickle diff --git a/make_datasets/electrode_positions.py b/make_datasets/electrode_positions.py index 09e2bb6..5be2ec4 100644 --- a/make_datasets/electrode_positions.py +++ b/make_datasets/electrode_positions.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_bci_ner.py b/make_datasets/make_bci_ner.py index e317d36..21f70c4 100644 --- a/make_datasets/make_bci_ner.py +++ b/make_datasets/make_bci_ner.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_boas.py b/make_datasets/make_boas.py index 19722d1..288bd94 100644 --- a/make_datasets/make_boas.py +++ b/make_datasets/make_boas.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the BOAS headband pre-training corpus into a pooled LMDB.""" import os diff --git a/make_datasets/make_chb_mit.py b/make_datasets/make_chb_mit.py index ca47d5c..4dd8c7f 100644 --- a/make_datasets/make_chb_mit.py +++ b/make_datasets/make_chb_mit.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the CHB-MIT seizure corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_gwd.py b/make_datasets/make_gwd.py index a44b957..f3b6999 100644 --- a/make_datasets/make_gwd.py +++ b/make_datasets/make_gwd.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the GWD pre-training corpus into a pooled LMDB.""" import os diff --git a/make_datasets/make_isruc.py b/make_datasets/make_isruc.py index 99daeee..8fac51c 100644 --- a/make_datasets/make_isruc.py +++ b/make_datasets/make_isruc.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_mental_arithmetic.py b/make_datasets/make_mental_arithmetic.py index 2096e81..d1f5f3a 100644 --- a/make_datasets/make_mental_arithmetic.py +++ b/make_datasets/make_mental_arithmetic.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the mental arithmetic corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_mumtaz.py b/make_datasets/make_mumtaz.py index 5d31bb2..44e587e 100644 --- a/make_datasets/make_mumtaz.py +++ b/make_datasets/make_mumtaz.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the Mumtaz depression corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_neonate.py b/make_datasets/make_neonate.py index d0ccff7..64014d8 100644 --- a/make_datasets/make_neonate.py +++ b/make_datasets/make_neonate.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the Helsinki neonatal seizure corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_physionet_mi.py b/make_datasets/make_physionet_mi.py index c32e763..69b6b08 100644 --- a/make_datasets/make_physionet_mi.py +++ b/make_datasets/make_physionet_mi.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the PhysioNet motor imagery corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_seed.py b/make_datasets/make_seed.py index 20879ac..bbb73e2 100644 --- a/make_datasets/make_seed.py +++ b/make_datasets/make_seed.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_seed_fra_ger.py b/make_datasets/make_seed_fra_ger.py index 09f38ec..4c146dc 100644 --- a/make_datasets/make_seed_fra_ger.py +++ b/make_datasets/make_seed_fra_ger.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_seed_v.py b/make_datasets/make_seed_v.py index 220d22e..0a23450 100644 --- a/make_datasets/make_seed_v.py +++ b/make_datasets/make_seed_v.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the SEED-V emotion corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_seed_vig.py b/make_datasets/make_seed_vig.py index 775c08a..7a5abf7 100644 --- a/make_datasets/make_seed_vig.py +++ b/make_datasets/make_seed_vig.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the SEED-VIG vigilance corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_shu_mi.py b/make_datasets/make_shu_mi.py index e4bfd71..cad3319 100644 --- a/make_datasets/make_shu_mi.py +++ b/make_datasets/make_shu_mi.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/make_datasets/make_sleepedfx.py b/make_datasets/make_sleepedfx.py index 7163337..07f2a2a 100644 --- a/make_datasets/make_sleepedfx.py +++ b/make_datasets/make_sleepedfx.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the Sleep-EDFx pre-training corpus into a pooled LMDB.""" import os diff --git a/make_datasets/make_stew.py b/make_datasets/make_stew.py index 509d173..91df147 100644 --- a/make_datasets/make_stew.py +++ b/make_datasets/make_stew.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the STEW workload corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_tuab.py b/make_datasets/make_tuab.py index 74aa6ee..f783240 100644 --- a/make_datasets/make_tuab.py +++ b/make_datasets/make_tuab.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + """Preprocess the TUAB abnormality-detection corpus into train/val/test LMDBs.""" import os diff --git a/make_datasets/make_tueg.py b/make_datasets/make_tueg.py index 7d5eff4..ab6de28 100644 --- a/make_datasets/make_tueg.py +++ b/make_datasets/make_tueg.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 diff --git a/models/model_heads/__init__.py b/models/model_heads/__init__.py new file mode 100644 index 0000000..1f9c938 --- /dev/null +++ b/models/model_heads/__init__.py @@ -0,0 +1,21 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 index dc9d9d1..5b19f6a 100644 --- a/models/model_heads/mlp_classification_head.py +++ b/models/model_heads/mlp_classification_head.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn from timm.layers import trunc_normal_ diff --git a/models/model_heads/mlp_regression_head.py b/models/model_heads/mlp_regression_head.py index aad7a56..fc983ac 100644 --- a/models/model_heads/mlp_regression_head.py +++ b/models/model_heads/mlp_regression_head.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn from timm.layers import trunc_normal_ diff --git a/models/model_heads/patch_reconstruction_head.py b/models/model_heads/patch_reconstruction_head.py index 591ab39..efa7bbf 100644 --- a/models/model_heads/patch_reconstruction_head.py +++ b/models/model_heads/patch_reconstruction_head.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn diff --git a/models/model_heads/sequence_classification_head.py b/models/model_heads/sequence_classification_head.py index e3270cc..d61eb8f 100644 --- a/models/model_heads/sequence_classification_head.py +++ b/models/model_heads/sequence_classification_head.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn import torch.nn.functional as F diff --git a/models/modules/attention.py b/models/modules/attention.py index 29cd600..a5eebff 100644 --- a/models/modules/attention.py +++ b/models/modules/attention.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Dict, Iterable, Optional, Tuple import torch diff --git a/models/modules/patching.py b/models/modules/patching.py index 3c783ab..3c0eacc 100644 --- a/models/modules/patching.py +++ b/models/modules/patching.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn from einops import rearrange diff --git a/models/modules/pos_embed.py b/models/modules/pos_embed.py index 09d233a..b1aa4ac 100644 --- a/models/modules/pos_embed.py +++ b/models/modules/pos_embed.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch import torch.nn as nn from timm.layers import trunc_normal_ diff --git a/models/s_cerebro.py b/models/s_cerebro.py index 1d26bc0..85e33aa 100644 --- a/models/s_cerebro.py +++ b/models/s_cerebro.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Iterable, Optional import torch 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 index b47e22e..98ce6ba 100644 --- a/tasks/classification_task.py +++ b/tasks/classification_task.py @@ -1,9 +1,34 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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, @@ -40,7 +65,7 @@ def split_checkpoint_state_dict(checkpoint: Dict[str, Any]) -> Dict[str, Dict[st return {"model": dict(checkpoint), "model_head": {}} -class ClassificationTask(pl.LightningModule): +class ClassificationTask(SafetensorsCheckpointMixin, pl.LightningModule): """Fine-tuning task for EEG classification. Supports two input layouts: @@ -93,6 +118,11 @@ def __init__( 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() @@ -154,6 +184,7 @@ def forward(self, x: torch.Tensor, channel_positions: torch.Tensor) -> torch.Ten 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: @@ -196,15 +227,15 @@ 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(batch, "train") + 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(batch, "val") + 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(batch, "test") + return self._shared_step(as_signal_batch(batch), "test") def on_train_epoch_end(self) -> None: """Log aggregated training metrics.""" diff --git a/tasks/mae_pretraining.py b/tasks/mae_pretraining.py index 818dd72..323b97a 100644 --- a/tasks/mae_pretraining.py +++ b/tasks/mae_pretraining.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + from typing import Any, Dict, Optional, Tuple import hydra @@ -5,6 +27,9 @@ 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 @@ -26,7 +51,7 @@ def extract_encoder_state_dict(checkpoint: Dict[str, Any]) -> Dict[str, torch.Te } -class MaskedAutoencoderPretrainingTask(pl.LightningModule): +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 @@ -48,6 +73,11 @@ def __init__(self, hparams, masking_ratio: float = 0.5): 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 @@ -66,6 +96,7 @@ def on_after_batch_transfer(self, batch: Dict[str, Any], dataloader_idx: int) -> 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] @@ -91,13 +122,13 @@ def _shared_step(self, batch: Dict[str, Any]) -> torch.Tensor: def training_step(self, batch: Dict[str, Any], batch_idx: int) -> torch.Tensor: """Compute and log the training reconstruction loss.""" - loss = self._shared_step(batch) + 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(batch) + 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 diff --git a/tasks/regression_task.py b/tasks/regression_task.py index f7e4a8f..cb3c10c 100644 --- a/tasks/regression_task.py +++ b/tasks/regression_task.py @@ -1,9 +1,34 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* 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 @@ -55,7 +80,7 @@ def compute(self) -> torch.Tensor: return rmse / torch.clamp(norm, min=1e-8) -class RegressionTask(pl.LightningModule): +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 @@ -77,6 +102,11 @@ def __init__(self, hparams, freeze_backbone: bool = False, layerwise_lr_decay: f 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) @@ -127,6 +157,7 @@ def forward(self, x: torch.Tensor, channel_positions: torch.Tensor) -> torch.Ten 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 @@ -160,15 +191,15 @@ 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(batch, "train") + 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(batch, "val") + 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(batch, "test") + return self._shared_step(as_signal_batch(batch), "test") def on_train_epoch_end(self) -> None: """Log aggregated training metrics.""" diff --git a/tests/model_tests/test_attention.py b/tests/model_tests/test_attention.py index fde1314..6d7281e 100644 --- a/tests/model_tests/test_attention.py +++ b/tests/model_tests/test_attention.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch from models.modules.attention import ( diff --git a/tests/model_tests/test_pipeline.py b/tests/model_tests/test_pipeline.py index edd147d..f583480 100644 --- a/tests/model_tests/test_pipeline.py +++ b/tests/model_tests/test_pipeline.py @@ -1,3 +1,25 @@ +#*----------------------------------------------------------------------------* +#* 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: Glenn Anta Bucagu * +#* Author: BioFoundation Contributors * +#* * +#* Imported from the S-CEReBrO reference implementation (TimeFM). * +#*----------------------------------------------------------------------------* + import torch from omegaconf import OmegaConf From 616111a7ddbbc0fe0d49f3d950259112ba74a007 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 30 Jul 2026 14:17:30 +0000 Subject: [PATCH 04/14] Register S-CEReBrO and document the encoder/head split Makes S-CEReBrO discoverable through the same metadata and documentation surface as the other five families. biofoundation/model_registry.py New "s-cerebro" entry: EEG, windowed alternating-attention Transformer, encoder target models.s_cerebro.SCerebroEncoder, four head targets, three size variants, and BatchRequirements(channel_coords=True). The key is the case-folded display name, matching every existing entry, and a new test pins that invariant so the convention cannot drift. docs/adr/0001-two-electrode-geometry-representations.md Records why channel_locations and channel_coords both exist as first-class fields, why neither is derived from the other, and what that costs. channel_coords is strictly richer, since a midpoint is recoverable from an electrode pair but not the reverse, so making it canonical was the alternative considered. It was rejected because it would have changed how the five existing families obtain geometry. The consequence is recorded plainly: a dataset prepared for one representation cannot feed a family expecting the other without an explicit conversion. Implicit conversion inside require_batch_fields is ruled out, because a model that silently received midpoints where it expected electrode pairs would train without error on quietly wrong geometry. docs/model/SCEReBrO.md Input assumptions, preprocessing, architecture, the alternating attention schedule, the SSL objective, the three downstream layouts, size variants, and runnable commands for pre-training, fine-tuning, dataset switching, sleep staging, regression and linear probing. README.md, models/README.md, docs/README.md, docs/CITATIONS.md Model-zoo row, family table row, documentation index entry, publication-status row and BibTeX. models/README.md now explains that the repository carries two model shapes and which protocols describe the second. CONTRIBUTING.md Adds the additive-change rule for biofoundation/, the requirement to record expensive-to-reverse decisions as an ADR, guidance on the encoder/head option when adding a model, and how to run the PyTorch-dependent tests under tests/model_tests. Across the branch, eight pre-existing files are modified: the four contract modules, three test modules and requirements.txt. No existing model, task, dataset, criterion, experiment or run_train.py is touched, and nothing is deleted. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- CONTRIBUTING.md | 14 +- README.md | 4 +- biofoundation/model_registry.py | 19 +++ config/experiment/SCEReBrO_finetune.yaml | 2 +- config/experiment/SCEReBrO_pretrain.yaml | 2 +- docs/CITATIONS.md | 15 ++ docs/README.md | 2 + ...-two-electrode-geometry-representations.md | 31 ++++ docs/model/SCEReBrO.md | 140 ++++++++++++++++++ models/README.md | 3 + tests/test_model_registry.py | 13 +- 11 files changed, 239 insertions(+), 6 deletions(-) create mode 100644 docs/adr/0001-two-electrode-geometry-representations.md create mode 100644 docs/model/SCEReBrO.md 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..6f7ad22 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.03118) / [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/model_registry.py b/biofoundation/model_registry.py index 4e0c88f..75bcdc3 100644 --- a/biofoundation/model_registry.py +++ b/biofoundation/model_registry.py @@ -115,6 +115,25 @@ class ModelSpec: ), 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.03118", + batch_requirements=BatchRequirements(channel_coords=True), + venue="arXiv preprint", + 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/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml index d9ae46c..0b4319b 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -22,7 +22,7 @@ tag: SCEReBrO_finetune model_size: tiny # Names the registry entry whose BatchRequirements the task enforces on every batch. -model_family: scerebro +model_family: s-cerebro # Fine-tuning dataset. Defaults to TUAB; switch corpus with a single override, e.g. # python -u run_train.py +experiment=SCEReBrO_finetune \ diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml index 336f942..0f6dbca 100644 --- a/config/experiment/SCEReBrO_pretrain.yaml +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -22,7 +22,7 @@ tag: SCEReBrO_pretrain model_size: tiny # Names the registry entry whose BatchRequirements the task enforces on every batch. -model_family: scerebro +model_family: s-cerebro gpus: -1 num_nodes: 1 diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md index 630f1ef..91294b4 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 | [arXiv preprint](https://arxiv.org/abs/2607.03118). | ## FEMBA @@ -85,4 +86,18 @@ Please cite the paper corresponding to the model used in your work. Published pr } ``` +## S-CEReBrO + +```bibtex +@misc{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}, + year={2026}, + eprint={2607.03118}, + archivePrefix={arXiv}, + primaryClass={eess.SP}, + url={https://arxiv.org/abs/2607.03118} +} +``` + The LuMamba and PanLUNA entries should be updated with their final DOI and page numbers once the EUSIPCO and AICAS 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..e9e312b --- /dev/null +++ b/docs/model/SCEReBrO.md @@ -0,0 +1,140 @@ +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. + +### 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 prepared dataset by overriding the corpus and its montage: + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune \ + dataset_root=$DATA_PATH/finetuning/CHB-MIT \ + model.num_channels=16 model_head.num_classes=2 +``` + +Sleep staging and regression swap the head, task and criterion together: + +```bash +# ISRUC, one label per 30 s epoch +python -u run_train.py +experiment=SCEReBrO_finetune \ + dataset_root=$DATA_PATH/finetuning/ISRUC dataset_kind=sequence \ + model.num_channels=6 /model_head=sequence_classification_head + +# SEED-VIG, continuous target +python -u run_train.py +experiment=SCEReBrO_finetune \ + dataset_root=$DATA_PATH/finetuning/SEED-VIG label_mode=regression \ + model.num_channels=17 /model_head=mlp_regression_head \ + /task=finetune_regression_task_SCEReBrO /criterion=mse_criterion +``` + +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. + +### Pretrained Weights + +Checkpoints are published at [PulpBio/S-CEReBrO](https://huggingface.co/PulpBio/S-CEReBrO). Encoder weights load independently of the head, and tensors whose shapes do not match the current model are skipped rather than forced, so an encoder pre-trained at 64 channels can seed a 22-channel fine-tune. The loader prints how many tensors were loaded, skipped, and unexpected, so a silent no-op load is visible in the logs. 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/tests/test_model_registry.py b/tests/test_model_registry.py index f7a0a26..0d1ebb4 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -30,7 +30,14 @@ 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) @@ -69,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") From 044374aaa0376d86de0b68bcd941a111971691c0 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 10:16:36 +0000 Subject: [PATCH 05/14] Fix S-CEReBrO composition, head overrides and docs; verify against main Verification with hydra-core and PyTorch actually installed, which was not possible when the previous commits were written. Three real defects surfaced and are fixed. Hydra defaults ordering Both S-CEReBrO experiments failed to compose. Hydra requires additions to the defaults list before overrides, so "- /model_head: ..." must precede the "- override /..." entries rather than follow them. Reordered in both experiments. Head-specific keys in the fine-tuning experiment SCEReBrO_finetune set num_classes and num_patches under model_head, which are MlpClassificationHead parameters. Selecting any other head then failed, because the classification keys were merged onto a head that does not accept them. The block is removed: each head config in config/model_head/ carries its own defaults, and per dataset values are given as model_head.num_classes=... on the command line. Documentation The head, task and criterion swap examples used a leading slash. That form belongs in a defaults list inside a config file; a command-line override takes the group name alone. Corrected, with a note explaining the distinction. arXiv identifier Updated to 2607.27913 in the registry, README and CITATIONS. The Hugging Face URL was already correct. tests/test_hydra_composition.py New test composing every file in config/experiment/, not only those a ModelSpec names. The defaults-ordering defect above produced a valid-looking YAML file that Hydra rejected only at run time, and nothing would have caught it. Experiments that do not compose are recorded in KNOWN_UNCOMPOSABLE_EXPERIMENTS with the reason, and the test fails if one starts composing, so an entry cannot rot. FEMBA_quantized is listed: it predates this branch, fails identically on main, is referenced by no registry entry, and has two independent faults - its defaults list requires scheduler/constant_lr, which does not exist, and its target misspells test_24_femba_full_expand2 as test_24_femba_full_expland2. Left unfixed here to keep this branch free of changes to the existing families. Verified All 10 pre-existing experiments resolve byte-identically to main. LUNA, TinyMyo and PanLUNA: identical parameter counts, state_dict signatures and forward-pass output hashes on main and on this branch, from deterministically filled weights. FEMBA and LuMamba cannot be imported without CUDA-built mamba_ssm, on main or here; their sources and resolved configs are unchanged, so their behaviour is unchanged. S-CEReBrO: 16 upstream unit tests pass, and the pre-training, classification, regression and sequence-classification tasks each complete a training step with a finite loss. require_batch_fields rejects a batch without channel_coords, which confirms the registry lookup is live at runtime. All six commands in docs/model/SCEReBrO.md compose. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- README.md | 2 +- biofoundation/model_registry.py | 2 +- config/experiment/SCEReBrO_finetune.yaml | 20 +++++++----- config/experiment/SCEReBrO_pretrain.yaml | 7 +++-- docs/CITATIONS.md | 6 ++-- docs/model/SCEReBrO.md | 14 ++++++--- tests/test_hydra_composition.py | 39 ++++++++++++++++++++++++ 7 files changed, 71 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index 6f7ad22..75be048 100644 --- a/README.md +++ b/README.md @@ -19,7 +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.03118) / [Hugging Face](https://huggingface.co/PulpBio/S-CEReBrO) | +| [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. diff --git a/biofoundation/model_registry.py b/biofoundation/model_registry.py index 75bcdc3..d5370cd 100644 --- a/biofoundation/model_registry.py +++ b/biofoundation/model_registry.py @@ -123,7 +123,7 @@ class ModelSpec: pretrain_experiment="SCEReBrO_pretrain", finetune_experiment="SCEReBrO_finetune", huggingface_url="https://huggingface.co/PulpBio/S-CEReBrO", - paper_url="https://arxiv.org/abs/2607.03118", + paper_url="https://arxiv.org/abs/2607.27913", batch_requirements=BatchRequirements(channel_coords=True), venue="arXiv preprint", head_targets=( diff --git a/config/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml index 0b4319b..27a3526 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -49,22 +49,28 @@ pretrained_checkpoint_path: null pretrained_safetensors_path: null 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. + # + # For SEED-VIG vigilance regression, change these together on the command line + # (no leading slash in CLI overrides): + # model_head=mlp_regression_head task=finetune_regression_task_SCEReBrO \ + # criterion=mse_criterion label_mode=regression + - /model_head: mlp_classification_head - override /data_module: finetune_data_module_SCEReBrO - override /model: SCEReBrO_tiny # or SCEReBrO_small, SCEReBrO_base - override /task: finetune_task_SCEReBrO - override /criterion: ce_criterion - override /scheduler: cosine - # For SEED-VIG vigilance regression, override these three together: - # /model_head: mlp_regression_head, /task: finetune_regression_task_SCEReBrO, - # /criterion: mse_criterion, and set label_mode=regression. - - /model_head: mlp_classification_head model: num_channels: 22 # TUAB montage -model_head: - num_classes: 2 - num_patches: 10 # 10 s window at one patch per second +# model_head values are not set here. Each head config in config/model_head/ carries +# its own defaults, so swapping the head group does not leave classification-only keys +# such as num_classes behind for a head that does not accept them. Override per +# dataset on the command line, e.g. model_head.num_classes=4 model_head.num_patches=4. task: freeze_backbone: False # True gives a linear-probe style run diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml index 0f6dbca..afc9bee 100644 --- a/config/experiment/SCEReBrO_pretrain.yaml +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -40,14 +40,15 @@ pretrained_checkpoint_path: null pretrained_safetensors_path: null 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_head is not part of the global defaults, so it is added here rather than - # overridden. The five bundled families compose exactly as before. - - /model_head: patch_reconstruction_head model: num_channels: 64 diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md index 91294b4..f037ef3 100644 --- a/docs/CITATIONS.md +++ b/docs/CITATIONS.md @@ -11,7 +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 | [arXiv preprint](https://arxiv.org/abs/2607.03118). | +| S-CEReBrO | [arXiv preprint](https://arxiv.org/abs/2607.27913). | ## FEMBA @@ -93,10 +93,10 @@ Please cite the paper corresponding to the model used in your work. Published pr 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}, year={2026}, - eprint={2607.03118}, + eprint={2607.27913}, archivePrefix={arXiv}, primaryClass={eess.SP}, - url={https://arxiv.org/abs/2607.03118} + url={https://arxiv.org/abs/2607.27913} } ``` diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md index e9e312b..1913e94 100644 --- a/docs/model/SCEReBrO.md +++ b/docs/model/SCEReBrO.md @@ -109,7 +109,7 @@ Fine-tune on another prepared dataset by overriding the corpus and its montage: ```bash python -u run_train.py +experiment=SCEReBrO_finetune \ dataset_root=$DATA_PATH/finetuning/CHB-MIT \ - model.num_channels=16 model_head.num_classes=2 + model.num_channels=16 model_head.num_classes=2 model_head.num_patches=10 ``` Sleep staging and regression swap the head, task and criterion together: @@ -118,15 +118,21 @@ Sleep staging and regression swap the head, task and criterion together: # ISRUC, one label per 30 s epoch python -u run_train.py +experiment=SCEReBrO_finetune \ dataset_root=$DATA_PATH/finetuning/ISRUC dataset_kind=sequence \ - model.num_channels=6 /model_head=sequence_classification_head + model.num_channels=6 model_head=sequence_classification_head # SEED-VIG, continuous target python -u run_train.py +experiment=SCEReBrO_finetune \ dataset_root=$DATA_PATH/finetuning/SEED-VIG label_mode=regression \ - model.num_channels=17 /model_head=mlp_regression_head \ - /task=finetune_regression_task_SCEReBrO /criterion=mse_criterion + model.num_channels=17 model_head=mlp_regression_head \ + task=finetune_regression_task_SCEReBrO criterion=mse_criterion ``` +A config group is selected on the command line without a leading slash +(`model_head=...`); the leading slash form is only used inside a defaults list in a +config file. Head parameters come from the selected head's own config, so override +them per dataset with `model_head.num_classes=4 model_head.num_patches=4` rather than +expecting the experiment to carry classification-specific keys. + A linear-probe style run freezes the encoder blocks while leaving tokenisation and the embeddings trainable: ```bash diff --git a/tests/test_hydra_composition.py b/tests/test_hydra_composition.py index 1e754c1..354f13b 100644 --- a/tests/test_hydra_composition.py +++ b/tests/test_hydra_composition.py @@ -33,6 +33,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 +71,33 @@ 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", + ) + if __name__ == "__main__": unittest.main() From a8c36546fa8fa79131bf53f7daf6eec86d08c2a7 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 10:35:02 +0000 Subject: [PATCH 06/14] Add synthetic-data generator and fix LMDB reopen during the test pass Running pre-training and fine-tuning end to end for the first time surfaced one real integration defect, fixed here, and produced the tooling to reproduce the run. make_datasets/make_dummy_scerebro_dataset.py Generates synthetic corpora in the exact on-disk formats the readers expect: pickled dictionaries plus a meta.json key index for LMDBDataset, and fixed-size raw blobs plus a key list for TUEGDataset, with padded channels written as all-zero coordinates so num_padded_channels is recovered as it is from real data. Covers the pre-training union, TUAB, ISRUC and SEED-VIG. A label is written only for corpora that have one. Pre-training samples must carry no label at all, because they are concatenated with TUEG, which has none, and a batch mixing samples with and without the key cannot be collated. datasets/lmdb_dataset.py Open the LMDB environment lazily on first read rather than in the constructor. run_train.py deletes and re-instantiates the data module before the rank-zero test pass, and LMDB refuses to open the same path twice in one process, so every fine-tuning run with final_test enabled aborted after training with "The environment ... is already open in this process". The upstream runner reuses one data module object and never hit this, so the defect only appears in BioFoundation. Lazy opening also stops an open handle being inherited by forked dataloader workers. This mirrors TUEGDataset, which already manages its environment this way. Where no meta.json index exists the key scan now uses a temporary environment that is closed again, so construction still leaves no handle open. This is the only vendored file that differs from upstream beyond its licence header. docs/model/SCEReBrO.md New smoke-test section with the generator invocation and the CPU overrides, and a note that on PyTorch 2.6 and newer the final validation and test passes fail when reloading a checkpoint, because torch.load now defaults to weights_only=True and the tasks store their Hydra configuration in the checkpoint. That affects run_train.py for every family; the workaround is TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1. Verified on CPU with generated data Pre-training: 9 corpora concatenated to 216 windows, one epoch, train and val loss logged, checkpoints written. Fine-tuning: encoder loaded from the pre-training checkpoint with 97 of 97 tensors matched and none skipped, two epochs, then validation and test with the full metric set reported. All 27 contract tests pass with and without hydra installed, the 16 vendored unit tests pass, and the 10 pre-existing experiments still resolve identically to main. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- datasets/lmdb_dataset.py | 30 ++- docs/model/SCEReBrO.md | 43 ++++ make_datasets/make_dummy_scerebro_dataset.py | 255 +++++++++++++++++++ 3 files changed, 323 insertions(+), 5 deletions(-) create mode 100644 make_datasets/make_dummy_scerebro_dataset.py diff --git a/datasets/lmdb_dataset.py b/datasets/lmdb_dataset.py index d62583f..ada3b04 100644 --- a/datasets/lmdb_dataset.py +++ b/datasets/lmdb_dataset.py @@ -132,14 +132,27 @@ def __init__( self.cache_size = cache_size self.lmdb_path, self.meta_path = resolve_lmdb_path(path, split) - self.env = lmdb.open( - self.lmdb_path, readonly=True, lock=False, readahead=True, max_readers=64, map_async=True - ) + 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: @@ -150,8 +163,14 @@ def _load_keys(self) -> List[bytes]: raise ValueError(f"Index file {self.meta_path} contains no keys") return [key.encode() if isinstance(key, str) else key for key in keys] - with self.env.begin(buffers=True) as txn: - return [bytes(key) for key in txn.cursor().iternext(keys=True, values=False)] + # 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.""" @@ -201,6 +220,7 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: 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) diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md index 1913e94..551fba2 100644 --- a/docs/model/SCEReBrO.md +++ b/docs/model/SCEReBrO.md @@ -141,6 +141,49 @@ 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 Checkpoints are published at [PulpBio/S-CEReBrO](https://huggingface.co/PulpBio/S-CEReBrO). Encoder weights load independently of the head, and tensors whose shapes do not match the current model are skipped rather than forced, so an encoder pre-trained at 64 channels can seed a 22-channel fine-tune. The loader prints how many tensors were loaded, skipped, and unexpected, so a silent no-op load is visible in the logs. diff --git a/make_datasets/make_dummy_scerebro_dataset.py b/make_datasets/make_dummy_scerebro_dataset.py new file mode 100644 index 0000000..f7bc65d --- /dev/null +++ b/make_datasets/make_dummy_scerebro_dataset.py @@ -0,0 +1,255 @@ +#*----------------------------------------------------------------------------* +#* 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 * +#*----------------------------------------------------------------------------* + +"""Generate synthetic S-CEReBrO datasets for smoke-testing the training pipeline. + +This writes data in exactly the on-disk formats the real readers expect, so a run +started against it exercises the same code path as a run against prepared data. The +signals are band-limited noise and the labels are random: the output is for checking +that a pipeline runs, not for measuring anything. + +Two formats are produced. + +``LMDBDataset`` corpora + An LMDB directory per split whose values are pickled dictionaries with ``eeg``, + ``channel_coords`` and optionally ``label``, plus a ``{split}.meta.json`` index + listing the keys so start-up does not scan the database. + +``TUEGDataset`` corpus + A single LMDB whose values are fixed-size raw blobs holding a padded + ``float32[max_channels, timesteps]`` waveform followed by + ``float32[max_channels, 2, 3]`` coordinates, plus a newline-separated key list. + Padded channels are written with all-zero coordinates, which is how the reader + recovers ``num_padded_channels``. + +Example: + python -m make_datasets.make_dummy_scerebro_dataset --output $DATA_PATH +""" + +import argparse +import json +import os +import pickle +from pathlib import Path +from typing import Optional, Sequence + +import lmdb +import numpy as np + +SAMPLING_RATE = 200 +MAP_SIZE = 1 << 30 + +# Pre-training corpora named by config/data_module/pretrain_data_module_SCEReBrO.yaml. +PRETRAIN_CORPORA = ("SEED", "SEED-IV", "SEED-GER", "SEED-FRA", "BOAS", "GWD", "SleepEDFx", "BCI-NER") + + +def make_channel_coords(num_channels: int, rng: np.random.Generator) -> np.ndarray: + """Return plausible unit-sphere electrode-pair coordinates of shape (C, 2, 3). + + Each channel is described by two electrodes. The first is placed on a sphere of + unit radius and the second is a nearby point, which is what a bipolar derivation + or a scalp-electrode-plus-reference channel looks like to the encoder. + """ + theta = rng.uniform(0, np.pi, size=num_channels) + phi = rng.uniform(0, 2 * np.pi, size=num_channels) + first = np.stack( + [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)], axis=-1 + ) + second = first + rng.normal(scale=0.15, size=(num_channels, 3)) + second /= np.linalg.norm(second, axis=-1, keepdims=True) + return np.stack([first, second], axis=1).astype(np.float32) + + +def make_eeg(num_channels: int, num_timesteps: int, rng: np.random.Generator) -> np.ndarray: + """Return band-limited noise of shape (C, T), loosely resembling scalp EEG.""" + t = np.arange(num_timesteps) / SAMPLING_RATE + signal = np.zeros((num_channels, num_timesteps), dtype=np.float32) + for freq in (2.0, 6.0, 10.0, 21.0): + amplitude = rng.uniform(0.2, 1.0, size=(num_channels, 1)) + phase = rng.uniform(0, 2 * np.pi, size=(num_channels, 1)) + signal += (amplitude * np.sin(2 * np.pi * freq * t + phase)).astype(np.float32) + signal += rng.normal(scale=0.3, size=signal.shape).astype(np.float32) + return signal * 20.0 + + +def write_lmdb_split( + directory: Path, + split: str, + num_samples: int, + num_channels: int, + num_timesteps: int, + rng: np.random.Generator, + label_kind: str = "none", + num_classes: int = 2, + sequence_length: Optional[int] = None, +) -> None: + """Write one ``{split}.lmdb`` of pickled samples plus its ``{split}.meta.json`` index. + + ``label_kind`` is one of ``none``, ``classification`` or ``regression``. It must be + ``none`` for pre-training corpora: those samples are concatenated with TUEG, which + carries no label, and a batch mixing samples with and without a ``label`` key + cannot be collated. + """ + if label_kind not in {"none", "classification", "regression"}: + raise ValueError(f"label_kind must be none, classification or regression, got {label_kind!r}") + directory.mkdir(parents=True, exist_ok=True) + path = directory / f"{split}.lmdb" + env = lmdb.open(str(path), map_size=MAP_SIZE) + keys = [] + + with env.begin(write=True) as txn: + for index in range(num_samples): + coords = make_channel_coords(num_channels, rng) + if sequence_length is None: + eeg = make_eeg(num_channels, num_timesteps, rng) + if label_kind == "classification": + label = int(rng.integers(num_classes)) + elif label_kind == "regression": + label = float(rng.uniform(0.0, 1.0)) + else: + label = None + else: + eeg = np.stack( + [make_eeg(num_channels, num_timesteps, rng) for _ in range(sequence_length)] + ) + label = ( + rng.integers(num_classes, size=sequence_length).astype(np.int64).tolist() + if label_kind == "classification" + else None + ) + + entry = {"eeg": eeg, "channel_coords": coords, "subject_id": f"S{index % 4:02d}"} + if label is not None: + entry["label"] = label + + key = f"{split}_{index:06d}".encode() + txn.put(key, pickle.dumps(entry)) + keys.append(key.decode()) + + env.close() + (directory / f"{split}.meta.json").write_text(json.dumps({"keys": keys})) + print(f" {path} ({num_samples} samples)") + + +def write_tueg( + directory: Path, + num_samples: int, + max_channels: int, + real_channels: int, + num_timesteps: int, + rng: np.random.Generator, +) -> None: + """Write the fixed-size-blob LMDB and key list that ``TUEGDataset`` reads.""" + directory.mkdir(parents=True, exist_ok=True) + path = directory / "TUEG.lmdb" + env = lmdb.open(str(path), map_size=MAP_SIZE) + keys = [] + + with env.begin(write=True) as txn: + for index in range(num_samples): + waveform = np.zeros((max_channels, num_timesteps), dtype=np.float32) + coords = np.zeros((max_channels, 2, 3), dtype=np.float32) + + # Vary how many channels are real so padding is genuinely exercised. + active = real_channels if index % 2 == 0 else max(4, real_channels // 2) + waveform[:active] = make_eeg(active, num_timesteps, rng) + coords[:active] = make_channel_coords(active, rng) + + key = f"tueg_{index:06d}".encode() + txn.put(key, waveform.tobytes() + coords.tobytes()) + keys.append(key.decode()) + + env.close() + (directory / "TUEG_keys.txt").write_text("\n".join(keys) + "\n") + print(f" {path} ({num_samples} windows, padded to {max_channels} channels)") + + +def main(argv: Optional[Sequence[str]] = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument( + "--output", required=True, + help="Root directory; pretraining/ and finetuning/ are created beneath it. " + "Use the same path you export as DATA_PATH.", + ) + parser.add_argument("--pretrain-samples", type=int, default=64, help="Windows per pre-training corpus.") + parser.add_argument("--finetune-samples", type=int, default=64, help="Windows in the fine-tuning train split.") + parser.add_argument("--max-channels", type=int, default=64, help="Channel capacity of the pre-training corpora.") + parser.add_argument("--slice-duration", type=int, default=30, help="Pre-training window length in seconds.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--datasets", nargs="+", default=["pretrain", "tuab"], + choices=["pretrain", "tuab", "isruc", "seed-vig"], + help="Which corpora to generate. 'pretrain' builds TUEG plus the eight LMDB corpora.", + ) + args = parser.parse_args(argv) + + rng = np.random.default_rng(args.seed) + root = Path(args.output).expanduser().resolve() + pretrain_timesteps = SAMPLING_RATE * args.slice_duration + + if "pretrain" in args.datasets: + print(f"Pre-training corpora -> {root/'pretraining'}") + write_tueg( + root / "pretraining" / "TUEG", args.pretrain_samples, + args.max_channels, args.max_channels // 2, pretrain_timesteps, rng, + ) + for corpus in PRETRAIN_CORPORA: + # Fewer channels than capacity, so the loader's zero-padding path runs. + write_lmdb_split( + root / "pretraining" / corpus, "train", args.pretrain_samples, + num_channels=args.max_channels // 2, num_timesteps=pretrain_timesteps, + rng=rng, label_kind="none", + ) + + if "tuab" in args.datasets: + print(f"TUAB (22 channels, 10 s, 2 classes) -> {root/'finetuning'/'TUAB'}") + for split, count in (("train", args.finetune_samples), + ("val", max(8, args.finetune_samples // 4)), + ("test", max(8, args.finetune_samples // 4))): + write_lmdb_split( + root / "finetuning" / "TUAB", split, count, + num_channels=22, num_timesteps=10 * SAMPLING_RATE, + rng=rng, label_kind="classification", num_classes=2, + ) + + if "isruc" in args.datasets: + print(f"ISRUC (6 channels, 30 s x 20 epochs, 5 classes) -> {root/'finetuning'/'ISRUC'}") + for split, count in (("train", args.finetune_samples // 4), + ("val", 8), ("test", 8)): + write_lmdb_split( + root / "finetuning" / "ISRUC", split, max(4, count), + num_channels=6, num_timesteps=30 * SAMPLING_RATE, + rng=rng, label_kind="classification", num_classes=5, sequence_length=20, + ) + + if "seed-vig" in args.datasets: + print(f"SEED-VIG (17 channels, 8 s, continuous) -> {root/'finetuning'/'SEED-VIG'}") + for split, count in (("train", args.finetune_samples), + ("val", 8), ("test", 8)): + write_lmdb_split( + root / "finetuning" / "SEED-VIG", split, count, + num_channels=17, num_timesteps=8 * SAMPLING_RATE, + rng=rng, label_kind="regression", + ) + + print(f"\nDone. Point DATA_PATH at {root}") + + +if __name__ == "__main__": + main() From 1045e05b0bf505b60c094cc26e85e69f2b71a633 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 14:08:03 +0000 Subject: [PATCH 07/14] Make S-CEReBrO checkpoints distributable as safetensors; pin min-max normalisation Testing the Hugging Face checkpoint path for the first time found it was broken. tasks/mae_pretraining.py The task aliased the encoder's mask and pad tokens onto itself. Assigning an nn.Parameter to a second module registers it a second time, so every checkpoint carried both "model.mask_token" and "mask_token" backed by one storage, 101 tensors where the model has 99. safetensors refuses shared storage, so util/ckpt_to_safetensor.py failed with "Some tensors share memory" and no S-CEReBrO checkpoint could be published in the format the releases use. The aliases are gone and both tokens are read through self.model. Upstream never hit this because it distributes .ckpt rather than safetensors. tests/model_tests/test_pipeline.py reads the token through the encoder to match. Normalisation S-CEReBrO uses per-channel min-max scaling to [-1, 1], applied by LMDBDataset via apply_minmax and written into TUEG offline by make_tueg, which uses the same formula. The quantile normalisation configured in config/defaults.yaml belongs to the tasks the bundled families use, and the S-CEReBrO tasks never read input_normalization, so it was already inert. Both experiments now set input_normalization.normalize to False explicitly, so a resolved config cannot be misread as enabling it. No behaviour changes. make_datasets/make_dummy_scerebro_dataset.py Synthetic TUEG windows are now min-max normalised at write time, as make_tueg does. TUEGDataset returns the stored waveform unchanged, so without this the TUEG windows sat on a different scale from the LMDB corpora they are concatenated with. docs/model/SCEReBrO.md Pretrained Weights rewritten to match the other families: snapshot_download from PulpBio/S-CEReBrO, then pretrained_safetensors_path with the local path. Documents pretrained_checkpoint_path for a local .ckpt, the conversion command, how to read the loader's counts, and that a size mismatch is silently skipped rather than raised so the printed counts are what to check. Verified on CPU Checkpoint now has 99 tensors with no shared storage; conversion to safetensors succeeds and yields 97 encoder plus 2 head tensors. Fine-tuning from that safetensors file reports "[load:model] loaded=97 shape_mismatch=0 unexpected=0 total_target=97", trains, validates and tests, with the classification head left freshly initialised. Resolved configs confirm input_normalization.normalize=False and apply_minmax=True for both experiments. 27 contract tests, 16 vendored unit tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- config/experiment/SCEReBrO_finetune.yaml | 8 ++++ config/experiment/SCEReBrO_pretrain.yaml | 8 ++++ docs/model/SCEReBrO.md | 46 +++++++++++++++++++- make_datasets/make_dummy_scerebro_dataset.py | 9 +++- tasks/mae_pretraining.py | 12 +++-- tests/model_tests/test_pipeline.py | 2 +- 6 files changed, 78 insertions(+), 7 deletions(-) diff --git a/config/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml index 27a3526..c9fb29b 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -48,6 +48,14 @@ 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. diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml index afc9bee..9217d50 100644 --- a/config/experiment/SCEReBrO_pretrain.yaml +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -39,6 +39,14 @@ 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. diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md index 551fba2..c7c6b3a 100644 --- a/docs/model/SCEReBrO.md +++ b/docs/model/SCEReBrO.md @@ -186,4 +186,48 @@ final_test=False`. ### Pretrained Weights -Checkpoints are published at [PulpBio/S-CEReBrO](https://huggingface.co/PulpBio/S-CEReBrO). Encoder weights load independently of the head, and tensors whose shapes do not match the current model are skipped rather than forced, so an encoder pre-trained at 64 channels can seed a 22-channel fine-tune. The loader prints how many tensors were loaded, skipped, and unexpected, so a silent no-op load is visible in the logs. +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. + +```python +from huggingface_hub import snapshot_download + +snapshot_download( + repo_id="PulpBio/S-CEReBrO", + local_dir="checkpoints/S-CEReBrO", +) +``` + +Select the matching model size and pass the local path. This is the same +`pretrained_safetensors_path` mechanism every other family uses; `run_train.py` reads +it before training starts and routes it to the task's safetensors loader. + +```bash +python -u run_train.py +experiment=SCEReBrO_finetune model=SCEReBrO_tiny \ + pretrained_safetensors_path=/absolute/path/to/checkpoints/S-CEReBrO/SCEReBrO_tiny.safetensors +``` + +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/make_dummy_scerebro_dataset.py b/make_datasets/make_dummy_scerebro_dataset.py index f7bc65d..ff4cedd 100644 --- a/make_datasets/make_dummy_scerebro_dataset.py +++ b/make_datasets/make_dummy_scerebro_dataset.py @@ -168,7 +168,14 @@ def write_tueg( # Vary how many channels are real so padding is genuinely exercised. active = real_channels if index % 2 == 0 else max(4, real_channels // 2) - waveform[:active] = make_eeg(active, num_timesteps, rng) + signal = make_eeg(active, num_timesteps, rng) + # TUEGDataset returns the stored waveform unchanged, so windows are + # min-max normalised offline exactly as make_tueg does. Without this the + # TUEG windows would be on a different scale from the LMDB corpora they + # are concatenated with. + minimum = signal.min(axis=1, keepdims=True) + maximum = signal.max(axis=1, keepdims=True) + waveform[:active] = ((signal - minimum) / (maximum - minimum + 1e-10) - 0.5) * 2.0 coords[:active] = make_channel_coords(active, rng) key = f"tueg_{index:06d}".encode() diff --git a/tasks/mae_pretraining.py b/tasks/mae_pretraining.py index 323b97a..469169c 100644 --- a/tasks/mae_pretraining.py +++ b/tasks/mae_pretraining.py @@ -82,8 +82,12 @@ def __init__(self, hparams, masking_ratio: float = 0.5): self.patch_size = self.hparams.model.patch_size self.num_channels = self.hparams.model.num_channels self.embed_dim = self.hparams.model.embed_dim - self.mask_token = self.model.mask_token - self.pad_token = self.model.pad_token + # 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]: @@ -157,7 +161,7 @@ def prepare_tokens( 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.pad_token.to(tokens.dtype), tokens) + 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 @@ -194,7 +198,7 @@ def mask_tokens( 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.mask_token.to(tokens.dtype), tokens) + 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]: diff --git a/tests/model_tests/test_pipeline.py b/tests/model_tests/test_pipeline.py index f583480..e304143 100644 --- a/tests/model_tests/test_pipeline.py +++ b/tests/model_tests/test_pipeline.py @@ -129,7 +129,7 @@ def test_token_mask_marks_exactly_the_replaced_tokens(): masked, token_mask = task.mask_tokens(tokens.clone(), attn_mask=None) - replaced = (masked == task.mask_token.detach()).all(dim=-1) + 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]) From f1f719189b589ec674fc882b3383a64d6ac514c4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 14:55:33 +0000 Subject: [PATCH 08/14] Tie model_size to the model group and document variable channel counts Answers three questions about using the published checkpoints, and removes a way the size label could disagree with the encoder it names. Channel counts The variable-montage mechanism transferred intact and is now documented and demonstrated. Nothing in the state dict depends on num_channels: the temporal position table is sized by max_timesteps and shared across channels, the channel embedding is an MLP over 3D electrode coordinates with no per-channel parameters, and the attention blocks take the channel count only as a reshape argument. Loading the 64-channel checkpoint into encoders built for 1, 4, 6, 16, 22, 32 and 64 channels gives loaded=97 shape_mismatch=0 unexpected=0 in every case. docs/model/SCEReBrO.md gains a Channel Counts section separating model.num_channels, which must equal what the dataset yields, from model.max_channels, which must stay at the pre-training value or the checkpoint will not match, and noting that the encoder raises on a mismatch rather than silently reshaping. model_size Moved from the experiments into config/model/SCEReBrO_{tiny,small,base}.yaml, so selecting a group sets the label. It previously sat in the experiment while the encoder came from the group, so model=SCEReBrO_small without model_size=small mislabelled the output directory, and now that model_size is interpolated into the checkpoint path it would have selected the wrong weights. Two contract tests pin this: a model config declaring model_size must agree with its own filename, and no experiment may restate a model_size its selected group already declares. Families that never interpolate it may still omit it entirely, and those that declare it only in the experiment are unaffected. Checkpoint location The fine-tuning experiment gains pretrained_root, defaulting to ${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO, so a checkpoint can be named without an absolute path: 'pretrained_safetensors_path=${pretrained_root}/SCEReBrO_${model_size}.safetensors' Switching size then needs one change and the checkpoint follows. The documented snapshot_download call writes to that directory. ${hydra:runtime.choices.model} would remove the model_size indirection entirely but does not resolve under hydra.compose, so it would break the composition tests; the group-declared label achieves the same guarantee without that cost. Verified on CPU Fine-tuning the 64-channel checkpoint onto a 6-channel ISRUC montage through the interpolated path: loaded=97 shape_mismatch=0 unexpected=0, trains, validates, tests. All three sizes resolve to their matching checkpoint filename. 29 contract tests, 16 vendored tests, 12 of 12 documented commands compose, and the 10 pre-existing experiments still resolve identically to main. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- config/experiment/SCEReBrO_finetune.yaml | 9 ++++- config/experiment/SCEReBrO_pretrain.yaml | 4 +- config/model/SCEReBrO_base.yaml | 5 +++ config/model/SCEReBrO_small.yaml | 5 +++ config/model/SCEReBrO_tiny.yaml | 5 +++ docs/model/SCEReBrO.md | 48 +++++++++++++++++++++--- tests/test_repository_contracts.py | 47 +++++++++++++++++++++++ 7 files changed, 116 insertions(+), 7 deletions(-) diff --git a/config/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml index c9fb29b..9b5191c 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -19,7 +19,14 @@ #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* tag: SCEReBrO_finetune -model_size: tiny + +# 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 diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml index 9217d50..c693caa 100644 --- a/config/experiment/SCEReBrO_pretrain.yaml +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -19,7 +19,9 @@ #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* tag: SCEReBrO_pretrain -model_size: tiny + +# model_size is supplied by the selected config/model group, so it always matches +# the encoder actually being built. # Names the registry entry whose BatchRequirements the task enforces on every batch. model_family: s-cerebro diff --git a/config/model/SCEReBrO_base.yaml b/config/model/SCEReBrO_base.yaml index f9e0394..9448918 100644 --- a/config/model/SCEReBrO_base.yaml +++ b/config/model/SCEReBrO_base.yaml @@ -18,6 +18,11 @@ #* Author: Glenn Anta Bucagu * #* 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 diff --git a/config/model/SCEReBrO_small.yaml b/config/model/SCEReBrO_small.yaml index 4a043b5..499eac4 100644 --- a/config/model/SCEReBrO_small.yaml +++ b/config/model/SCEReBrO_small.yaml @@ -18,6 +18,11 @@ #* Author: Glenn Anta Bucagu * #* 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 diff --git a/config/model/SCEReBrO_tiny.yaml b/config/model/SCEReBrO_tiny.yaml index 49f986c..15cc7f5 100644 --- a/config/model/SCEReBrO_tiny.yaml +++ b/config/model/SCEReBrO_tiny.yaml @@ -18,6 +18,11 @@ #* Author: Glenn Anta Bucagu * #* 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 diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md index c7c6b3a..960afeb 100644 --- a/docs/model/SCEReBrO.md +++ b/docs/model/SCEReBrO.md @@ -21,6 +21,26 @@ The encoder requires `channel_coords` of shape `(batch, channels, 2, 3)`: the 3D 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. @@ -188,24 +208,42 @@ final_test=False`. 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="checkpoints/S-CEReBrO", + local_dir=os.path.join(os.environ["CHECKPOINT_DIR"], "pretrained", "S-CEReBrO"), ) ``` -Select the matching model size and pass the local path. This is the same -`pretrained_safetensors_path` mechanism every other family uses; `run_train.py` reads -it before training starts and routes it to the task's safetensors loader. +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=/absolute/path/to/checkpoints/S-CEReBrO/SCEReBrO_tiny.safetensors + '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 diff --git a/tests/test_repository_contracts.py b/tests/test_repository_contracts.py index 6a383e5..be9d672 100644 --- a/tests/test_repository_contracts.py +++ b/tests/test_repository_contracts.py @@ -106,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 = ( From fc0ffef94acf291ab899669d8f72ef3cb7fe190a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 15:15:54 +0000 Subject: [PATCH 09/14] Freeze pre-training-only tokens so DDP fine-tuning does not hang Fine-tuning with the experiment's own defaults hangs at the first training step on a multi-GPU machine, with no error and no output. The encoder's mask_token and pad_token exist for masked pre-training. A fine-tuning forward pass never substitutes either, so neither receives a gradient. The fine-tuning experiment sets strategy: ddp and find_unused_parameters: False, and DistributedDataParallel's reducer waits for a gradient from every parameter it tracks, so the step never completes. Nothing is printed, which is why this reads as the run doing nothing rather than as a failure. This was not caught earlier because every smoke test forced trainer.strategy=auto and trainer.devices=1, which takes DDP out of the picture entirely. tasks/classification_task.py, tasks/regression_task.py New freeze_pretraining_only_parameters sets requires_grad=False on the encoder's mask and pad tokens, called from both fine-tuning tasks. That removes them from the set DDP tracks, which is cheaper than enabling find_unused_parameters and walking the autograd graph on every step. configure_optimizers already skips parameters that do not require grad, so no parameter group changes. LUNA disables its mask token for classification for the same reason; see models/LUNA.py. tests/model_tests/test_pipeline.py New test asserting that a fine-tuning backward leaves no trainable parameter without a gradient, so a future parameter that the fine-tuning path does not touch fails the suite rather than silently hanging a distributed run. Verified Both fine-tuning tasks now report zero trainable parameters without a gradient, where classification and regression each previously reported model.mask_token and model.pad_token. Fine-tuning from safetensors still trains, validates and tests with loaded=97 shape_mismatch=0 unexpected=0. 29 contract tests, 17 vendored tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- tasks/classification_task.py | 28 ++++++++++++++++++++++++++++ tasks/regression_task.py | 7 ++++++- tests/model_tests/test_pipeline.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 1 deletion(-) diff --git a/tasks/classification_task.py b/tasks/classification_task.py index 98ce6ba..fb15ef8 100644 --- a/tasks/classification_task.py +++ b/tasks/classification_task.py @@ -65,6 +65,32 @@ def split_checkpoint_state_dict(checkpoint: Dict[str, Any]) -> Dict[str, Dict[st 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. @@ -127,6 +153,8 @@ def __init__( 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() diff --git a/tasks/regression_task.py b/tasks/regression_task.py index cb3c10c..5fa0ba0 100644 --- a/tasks/regression_task.py +++ b/tasks/regression_task.py @@ -32,7 +32,10 @@ from torchmetrics import Metric from torchmetrics.regression import MeanSquaredError, PearsonCorrCoef, R2Score -from tasks.classification_task import split_checkpoint_state_dict +from tasks.classification_task import ( + freeze_pretraining_only_parameters, + split_checkpoint_state_dict, +) class NormalizedRootMeanSquaredError(Metric): @@ -116,6 +119,8 @@ def __init__(self, hparams, freeze_backbone: bool = False, layerwise_lr_decay: f 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() diff --git a/tests/model_tests/test_pipeline.py b/tests/model_tests/test_pipeline.py index e304143..8f22f93 100644 --- a/tests/model_tests/test_pipeline.py +++ b/tests/model_tests/test_pipeline.py @@ -173,3 +173,33 @@ def test_classification_head_pooling_modes_agree_on_output_shape(): 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}" From 4e61871254d7592c172728780090631c624c5620 Mon Sep 17 00:00:00 2001 From: g12bftd Date: Fri, 31 Jul 2026 17:49:09 +0200 Subject: [PATCH 10/14] Minor fixes --- biofoundation/model_registry.py | 2 +- config/experiment/SCEReBrO_finetune.yaml | 3 +-- config/experiment/SCEReBrO_pretrain.yaml | 6 ++---- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/biofoundation/model_registry.py b/biofoundation/model_registry.py index d5370cd..c37bcc0 100644 --- a/biofoundation/model_registry.py +++ b/biofoundation/model_registry.py @@ -125,7 +125,7 @@ class ModelSpec: huggingface_url="https://huggingface.co/PulpBio/S-CEReBrO", paper_url="https://arxiv.org/abs/2607.27913", batch_requirements=BatchRequirements(channel_coords=True), - venue="arXiv preprint", + venue="MICCAI 2026", head_targets=( "models.model_heads.patch_reconstruction_head.PatchReconstructionHead", "models.model_heads.mlp_classification_head.MlpClassificationHead", diff --git a/config/experiment/SCEReBrO_finetune.yaml b/config/experiment/SCEReBrO_finetune.yaml index c9fb29b..399fee7 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* tag: SCEReBrO_finetune @@ -43,7 +42,7 @@ training: True resume: False final_validate: True final_test: True -find_unused_parameters: False +find_unused_parameters: True pretrained_checkpoint_path: null pretrained_safetensors_path: null diff --git a/config/experiment/SCEReBrO_pretrain.yaml b/config/experiment/SCEReBrO_pretrain.yaml index 9217d50..1a8939f 100644 --- a/config/experiment/SCEReBrO_pretrain.yaml +++ b/config/experiment/SCEReBrO_pretrain.yaml @@ -15,13 +15,11 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* tag: SCEReBrO_pretrain model_size: tiny -# Names the registry entry whose BatchRequirements the task enforces on every batch. model_family: s-cerebro gpus: -1 @@ -92,7 +90,7 @@ trainer: num_nodes: ${num_nodes} devices: ${gpus} strategy: ddp - max_epochs: 50 + max_epochs: 100 check_val_every_n_epoch: 5 accumulate_grad_batches: 8 @@ -104,6 +102,6 @@ optimizer: scheduler: trainer: ${trainer} - warmup_epochs: 3 + warmup_epochs: 5 min_lr: 2.5e-7 warmup_lr_init: 2.5e-7 From 3b587ed7607d19611af3093ed901bed167176084 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 16:07:02 +0000 Subject: [PATCH 11/14] Select a fine-tuning corpus as one config group instead of five overrides Switching S-CEReBrO to a different fine-tuning corpus previously meant keeping up to five overrides consistent by hand: dataset_root, dataset_kind, label_mode, model.num_channels, and the model_head, task and criterion trio. Getting the trio partly right is not an error, it is a silently different run, and pairing a classification head's num_classes with a regression head is a crash. config/dataset/ New group with one file per prepared corpus: tuab, chb-mit, neonate, physionet-mi, shu-mi, stew, mumtaz, mental-arithmetic, seed-v, isruc and seed-vig. Each owns its corpus path, sample layout, label kind, channel count and the matching head, task and criterion, so the whole combination moves together: python -u run_train.py +experiment=SCEReBrO_finetune dataset=isruc Each file sets only the keys its selected head accepts, so seed-vig carries neither num_classes nor num_patches. config/experiment/SCEReBrO_finetune.yaml Gives up ownership of everything the group now holds. This is required rather than tidy: 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. That is exactly what an earlier probe showed, where a dataset group's head selection applied but its task, criterion and channel count did not. tests/test_hydra_composition.py Two tests. The first composes every dataset option and checks the head, task and label_mode agree on whether the run is regression, and that a head only receives keys its constructor accepts. The second fails if the experiment restates anything the group owns. Both were confirmed non-vacuous by reintroducing the corresponding mistake; the second initially missed a nested model.num_channels and its pattern was corrected. Non-S-CEReBrO models are unaffected. config/defaults.yaml is untouched, so the new group is reachable only from the S-CEReBrO fine-tuning experiment, and no existing experiment, model, task, dataset or criterion file is modified anywhere on this branch. All 10 pre-existing experiments still resolve byte-identically to main. Verified on CPU All 11 dataset options compose with the correct head, task and criterion. Real runs for dataset=tuab, dataset=isruc and dataset=seed-vig each load the pre-trained encoder with loaded=97 shape_mismatch=0 and complete training, validation and test, reporting accuracy for the two classification corpora and rmse, nrmse, r2 and pearson for the regression corpus. 31 contract tests pass with and without hydra installed, 17 vendored unit tests pass, and all 13 commands in docs/model/SCEReBrO.md compose. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- config/dataset/chb-mit.yaml | 38 +++++++++++++ config/dataset/isruc.yaml | 39 +++++++++++++ config/dataset/mental-arithmetic.yaml | 38 +++++++++++++ config/dataset/mumtaz.yaml | 38 +++++++++++++ config/dataset/neonate.yaml | 38 +++++++++++++ config/dataset/physionet-mi.yaml | 38 +++++++++++++ config/dataset/seed-v.yaml | 38 +++++++++++++ config/dataset/seed-vig.yaml | 37 +++++++++++++ config/dataset/shu-mi.yaml | 38 +++++++++++++ config/dataset/stew.yaml | 38 +++++++++++++ config/dataset/tuab.yaml | 38 +++++++++++++ config/experiment/SCEReBrO_finetune.yaml | 39 +++++-------- docs/model/SCEReBrO.md | 43 ++++++++------- tests/test_hydra_composition.py | 70 ++++++++++++++++++++++++ 14 files changed, 524 insertions(+), 46 deletions(-) create mode 100644 config/dataset/chb-mit.yaml create mode 100644 config/dataset/isruc.yaml create mode 100644 config/dataset/mental-arithmetic.yaml create mode 100644 config/dataset/mumtaz.yaml create mode 100644 config/dataset/neonate.yaml create mode 100644 config/dataset/physionet-mi.yaml create mode 100644 config/dataset/seed-v.yaml create mode 100644 config/dataset/seed-vig.yaml create mode 100644 config/dataset/shu-mi.yaml create mode 100644 config/dataset/stew.yaml create mode 100644 config/dataset/tuab.yaml diff --git a/config/dataset/chb-mit.yaml b/config/dataset/chb-mit.yaml new file mode 100644 index 0000000..50e4c8f --- /dev/null +++ b/config/dataset/chb-mit.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: Glenn Anta Bucagu * +#* 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..e8fd746 --- /dev/null +++ b/config/dataset/isruc.yaml @@ -0,0 +1,39 @@ +# @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: Glenn Anta Bucagu * +#* 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..32fc8b8 --- /dev/null +++ b/config/dataset/mental-arithmetic.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: Glenn Anta Bucagu * +#* 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..3881c10 --- /dev/null +++ b/config/dataset/mumtaz.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: Glenn Anta Bucagu * +#* 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..105642a --- /dev/null +++ b/config/dataset/neonate.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: Glenn Anta Bucagu * +#* 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..9a4fbf7 --- /dev/null +++ b/config/dataset/physionet-mi.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: Glenn Anta Bucagu * +#* 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..3becbb7 --- /dev/null +++ b/config/dataset/seed-v.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: Glenn Anta Bucagu * +#* 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..d8b11b2 --- /dev/null +++ b/config/dataset/seed-vig.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: Glenn Anta Bucagu * +#* 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..992a2b2 --- /dev/null +++ b/config/dataset/shu-mi.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: Glenn Anta Bucagu * +#* 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..5008f0c --- /dev/null +++ b/config/dataset/stew.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: Glenn Anta Bucagu * +#* 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..747b718 --- /dev/null +++ b/config/dataset/tuab.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: Glenn Anta Bucagu * +#* 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 index 9b5191c..05a493c 100644 --- a/config/experiment/SCEReBrO_finetune.yaml +++ b/config/experiment/SCEReBrO_finetune.yaml @@ -31,14 +31,11 @@ pretrained_root: ${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO # Names the registry entry whose BatchRequirements the task enforces on every batch. model_family: s-cerebro -# Fine-tuning dataset. Defaults to TUAB; switch corpus with a single override, e.g. -# python -u run_train.py +experiment=SCEReBrO_finetune \ -# dataset_root=${env:DATA_PATH}/finetuning/CHB-MIT model.num_channels=16 -# docs/model/SCEReBrO.md lists the channel count, window length and class count for -# each prepared dataset. -dataset_root: ${env:DATA_PATH}/finetuning/TUAB -dataset_kind: window # 'window', or 'sequence' for ISRUC sleep staging -label_mode: classification +# 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 @@ -64,29 +61,19 @@ 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. + # 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. # - # For SEED-VIG vigilance regression, change these together on the command line - # (no leading slash in CLI overrides): - # model_head=mlp_regression_head task=finetune_regression_task_SCEReBrO \ - # criterion=mse_criterion label_mode=regression - - /model_head: mlp_classification_head + # 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 /task: finetune_task_SCEReBrO - - override /criterion: ce_criterion - override /scheduler: cosine -model: - num_channels: 22 # TUAB montage - -# model_head values are not set here. Each head config in config/model_head/ carries -# its own defaults, so swapping the head group does not leave classification-only keys -# such as num_classes behind for a head that does not accept them. Override per -# dataset on the command line, e.g. model_head.num_classes=4 model_head.num_patches=4. - task: freeze_backbone: False # True gives a linear-probe style run layerwise_lr_decay: 0.95 diff --git a/docs/model/SCEReBrO.md b/docs/model/SCEReBrO.md index 960afeb..05f6393 100644 --- a/docs/model/SCEReBrO.md +++ b/docs/model/SCEReBrO.md @@ -124,34 +124,37 @@ python -u run_train.py +experiment=SCEReBrO_finetune \ pretrained_safetensors_path=/absolute/path/to/SCEReBrO_tiny.safetensors ``` -Fine-tune on another prepared dataset by overriding the corpus and its montage: +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_root=$DATA_PATH/finetuning/CHB-MIT \ - model.num_channels=16 model_head.num_classes=2 model_head.num_patches=10 +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 ``` -Sleep staging and regression swap the head, task and criterion together: +| `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 -# ISRUC, one label per 30 s epoch -python -u run_train.py +experiment=SCEReBrO_finetune \ - dataset_root=$DATA_PATH/finetuning/ISRUC dataset_kind=sequence \ - model.num_channels=6 model_head=sequence_classification_head - -# SEED-VIG, continuous target -python -u run_train.py +experiment=SCEReBrO_finetune \ - dataset_root=$DATA_PATH/finetuning/SEED-VIG label_mode=regression \ - model.num_channels=17 model_head=mlp_regression_head \ - task=finetune_regression_task_SCEReBrO criterion=mse_criterion +python -u run_train.py +experiment=SCEReBrO_finetune dataset=tuab model.num_channels=19 ``` -A config group is selected on the command line without a leading slash -(`model_head=...`); the leading slash form is only used inside a defaults list in a -config file. Head parameters come from the selected head's own config, so override -them per dataset with `model_head.num_classes=4 model_head.num_patches=4` rather than -expecting the experiment to carry classification-specific keys. +**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: diff --git a/tests/test_hydra_composition.py b/tests/test_hydra_composition.py index 354f13b..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 @@ -99,5 +100,74 @@ def test_every_experiment_file_composes_or_is_a_known_failure(self): ) + 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() From c71d7212b7f794275e9a0a3cd1907a160541ab14 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 16:09:32 +0000 Subject: [PATCH 12/14] Align CITATIONS with the MICCAI 2026 venue The registry entry for S-CEReBrO now records venue="MICCAI 2026", which the registry-driven citation test checks against docs/CITATIONS.md. The publication-status row and the BibTeX entry are updated to match: an inproceedings entry noting the proceedings are forthcoming, following the LuMamba and PanLUNA pattern, and the closing note now lists MICCAI alongside EUSIPCO and AICAS. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- docs/CITATIONS.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/CITATIONS.md b/docs/CITATIONS.md index f037ef3..ebe11ac 100644 --- a/docs/CITATIONS.md +++ b/docs/CITATIONS.md @@ -11,7 +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 | [arXiv preprint](https://arxiv.org/abs/2607.27913). | +| S-CEReBrO | Accepted at MICCAI 2026; proceedings forthcoming ([preprint](https://arxiv.org/abs/2607.27913)). | ## FEMBA @@ -89,10 +89,12 @@ Please cite the paper corresponding to the model used in your work. Published pr ## S-CEReBrO ```bibtex -@misc{bucagu2026scerebro, +@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}, @@ -100,4 +102,4 @@ 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. +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. From ec82da7941fd03cda4a62613817607acd6c1afcd Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 16:14:19 +0000 Subject: [PATCH 13/14] Remove the personal author line from file headers Every file carrying "Author: Glenn Anta Bucagu" now credits only "BioFoundation Contributors", matching the two experiment configs already corrected on this branch. 69 files, one deleted line each; no other content changed. No file that exists on main carried the line, so nothing pre-existing is touched. The line had been added to files created on this branch, including ones in shared directories such as criterion/, tasks/, models/modules/, datasets/ and make_datasets/, where it read as a claim over generic infrastructure. The BibTeX entry in docs/CITATIONS.md keeps "Bucagu, Glenn Anta" in its author list. That is the paper's authorship, not a file-header attribution. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha --- config/criterion/ce_criterion.yaml | 1 - config/criterion/focal_criterion.yaml | 1 - config/criterion/masked_reconstruction_loss.yaml | 1 - config/criterion/mse_criterion.yaml | 1 - config/data_module/finetune_data_module_SCEReBrO.yaml | 1 - config/data_module/pretrain_data_module_SCEReBrO.yaml | 1 - config/dataset/chb-mit.yaml | 1 - config/dataset/isruc.yaml | 1 - config/dataset/mental-arithmetic.yaml | 1 - config/dataset/mumtaz.yaml | 1 - config/dataset/neonate.yaml | 1 - config/dataset/physionet-mi.yaml | 1 - config/dataset/seed-v.yaml | 1 - config/dataset/seed-vig.yaml | 1 - config/dataset/shu-mi.yaml | 1 - config/dataset/stew.yaml | 1 - config/dataset/tuab.yaml | 1 - config/model/SCEReBrO_base.yaml | 1 - config/model/SCEReBrO_small.yaml | 1 - config/model/SCEReBrO_tiny.yaml | 1 - config/model_head/mlp_classification_head.yaml | 1 - config/model_head/mlp_regression_head.yaml | 1 - config/model_head/patch_reconstruction_head.yaml | 1 - config/model_head/sequence_classification_head.yaml | 1 - config/task/finetune_regression_task_SCEReBrO.yaml | 1 - config/task/finetune_task_SCEReBrO.yaml | 1 - config/task/pretrain_task_SCEReBrO.yaml | 1 - criterion/ce_criterion.py | 1 - criterion/focal_criterion.py | 1 - criterion/masked_reconstruction_loss.py | 1 - criterion/mse_criterion.py | 1 - data_module/finetuning_data_module.py | 1 - data_module/pretraining_data_module.py | 1 - datasets/lmdb_dataset.py | 1 - datasets/tueg_dataset.py | 1 - make_datasets/common.py | 1 - make_datasets/electrode_positions.py | 1 - make_datasets/make_bci_ner.py | 1 - make_datasets/make_boas.py | 1 - make_datasets/make_chb_mit.py | 1 - make_datasets/make_gwd.py | 1 - make_datasets/make_isruc.py | 1 - make_datasets/make_mental_arithmetic.py | 1 - make_datasets/make_mumtaz.py | 1 - make_datasets/make_neonate.py | 1 - make_datasets/make_physionet_mi.py | 1 - make_datasets/make_seed.py | 1 - make_datasets/make_seed_fra_ger.py | 1 - make_datasets/make_seed_v.py | 1 - make_datasets/make_seed_vig.py | 1 - make_datasets/make_shu_mi.py | 1 - make_datasets/make_sleepedfx.py | 1 - make_datasets/make_stew.py | 1 - make_datasets/make_tuab.py | 1 - make_datasets/make_tueg.py | 1 - models/model_heads/__init__.py | 1 - models/model_heads/mlp_classification_head.py | 1 - models/model_heads/mlp_regression_head.py | 1 - models/model_heads/patch_reconstruction_head.py | 1 - models/model_heads/sequence_classification_head.py | 1 - models/modules/attention.py | 1 - models/modules/patching.py | 1 - models/modules/pos_embed.py | 1 - models/s_cerebro.py | 1 - tasks/classification_task.py | 1 - tasks/mae_pretraining.py | 1 - tasks/regression_task.py | 1 - tests/model_tests/test_attention.py | 1 - tests/model_tests/test_pipeline.py | 1 - 69 files changed, 69 deletions(-) diff --git a/config/criterion/ce_criterion.yaml b/config/criterion/ce_criterion.yaml index c89eb81..0694a98 100644 --- a/config/criterion/ce_criterion.yaml +++ b/config/criterion/ce_criterion.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* criterion: diff --git a/config/criterion/focal_criterion.yaml b/config/criterion/focal_criterion.yaml index c5d7cf4..73beb5c 100644 --- a/config/criterion/focal_criterion.yaml +++ b/config/criterion/focal_criterion.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* criterion: diff --git a/config/criterion/masked_reconstruction_loss.yaml b/config/criterion/masked_reconstruction_loss.yaml index 9324c49..3e431f5 100644 --- a/config/criterion/masked_reconstruction_loss.yaml +++ b/config/criterion/masked_reconstruction_loss.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* criterion: diff --git a/config/criterion/mse_criterion.yaml b/config/criterion/mse_criterion.yaml index e288140..25b65ac 100644 --- a/config/criterion/mse_criterion.yaml +++ b/config/criterion/mse_criterion.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* criterion: diff --git a/config/data_module/finetune_data_module_SCEReBrO.yaml b/config/data_module/finetune_data_module_SCEReBrO.yaml index 5f8d6de..a4318a7 100644 --- a/config/data_module/finetune_data_module_SCEReBrO.yaml +++ b/config/data_module/finetune_data_module_SCEReBrO.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* data_module: diff --git a/config/data_module/pretrain_data_module_SCEReBrO.yaml b/config/data_module/pretrain_data_module_SCEReBrO.yaml index 5756a98..6e071ff 100644 --- a/config/data_module/pretrain_data_module_SCEReBrO.yaml +++ b/config/data_module/pretrain_data_module_SCEReBrO.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* data_module: diff --git a/config/dataset/chb-mit.yaml b/config/dataset/chb-mit.yaml index 50e4c8f..047a542 100644 --- a/config/dataset/chb-mit.yaml +++ b/config/dataset/chb-mit.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # CHB-MIT: seizure detection. 16 channels, 10 s windows, 2 classes. diff --git a/config/dataset/isruc.yaml b/config/dataset/isruc.yaml index e8fd746..31cf08d 100644 --- a/config/dataset/isruc.yaml +++ b/config/dataset/isruc.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # ISRUC: sleep staging. 6 channels, sequences of 20 consecutive 30 s epochs, 5 classes diff --git a/config/dataset/mental-arithmetic.yaml b/config/dataset/mental-arithmetic.yaml index 32fc8b8..3a0fac8 100644 --- a/config/dataset/mental-arithmetic.yaml +++ b/config/dataset/mental-arithmetic.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # MentalArithmetic: mental arithmetic. 20 channels, 5 s windows, 2 classes. diff --git a/config/dataset/mumtaz.yaml b/config/dataset/mumtaz.yaml index 3881c10..21af2b9 100644 --- a/config/dataset/mumtaz.yaml +++ b/config/dataset/mumtaz.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # Mumtaz: depression. 20 channels, 5 s windows, 2 classes. diff --git a/config/dataset/neonate.yaml b/config/dataset/neonate.yaml index 105642a..eba38e4 100644 --- a/config/dataset/neonate.yaml +++ b/config/dataset/neonate.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # Neonate: neonatal seizure detection. 18 channels, 5 s windows, 2 classes. diff --git a/config/dataset/physionet-mi.yaml b/config/dataset/physionet-mi.yaml index 9a4fbf7..31c09e2 100644 --- a/config/dataset/physionet-mi.yaml +++ b/config/dataset/physionet-mi.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # PhysioNet-MI: motor imagery. 64 channels, 4 s windows, 4 classes. diff --git a/config/dataset/seed-v.yaml b/config/dataset/seed-v.yaml index 3becbb7..3cd31bb 100644 --- a/config/dataset/seed-v.yaml +++ b/config/dataset/seed-v.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # SEED-V: emotion. 62 channels, 4 s windows, 5 classes. diff --git a/config/dataset/seed-vig.yaml b/config/dataset/seed-vig.yaml index d8b11b2..4bc322f 100644 --- a/config/dataset/seed-vig.yaml +++ b/config/dataset/seed-vig.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # SEED-VIG: vigilance regression. 17 channels, 8 s windows, continuous target in [0, 1]. diff --git a/config/dataset/shu-mi.yaml b/config/dataset/shu-mi.yaml index 992a2b2..880ba22 100644 --- a/config/dataset/shu-mi.yaml +++ b/config/dataset/shu-mi.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # SHU-MI: motor imagery. 32 channels, 4 s windows, 2 classes. diff --git a/config/dataset/stew.yaml b/config/dataset/stew.yaml index 5008f0c..ca68496 100644 --- a/config/dataset/stew.yaml +++ b/config/dataset/stew.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # STEW: workload. 14 channels, 4 s windows, 3 classes. diff --git a/config/dataset/tuab.yaml b/config/dataset/tuab.yaml index 747b718..c792865 100644 --- a/config/dataset/tuab.yaml +++ b/config/dataset/tuab.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # TUAB: abnormal vs normal. 22 channels, 10 s windows, 2 classes. diff --git a/config/model/SCEReBrO_base.yaml b/config/model/SCEReBrO_base.yaml index 9448918..c570d64 100644 --- a/config/model/SCEReBrO_base.yaml +++ b/config/model/SCEReBrO_base.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # Declared here rather than in the experiment so that selecting this group with diff --git a/config/model/SCEReBrO_small.yaml b/config/model/SCEReBrO_small.yaml index 499eac4..3adacfb 100644 --- a/config/model/SCEReBrO_small.yaml +++ b/config/model/SCEReBrO_small.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # Declared here rather than in the experiment so that selecting this group with diff --git a/config/model/SCEReBrO_tiny.yaml b/config/model/SCEReBrO_tiny.yaml index 15cc7f5..26b7757 100644 --- a/config/model/SCEReBrO_tiny.yaml +++ b/config/model/SCEReBrO_tiny.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* # Declared here rather than in the experiment so that selecting this group with diff --git a/config/model_head/mlp_classification_head.yaml b/config/model_head/mlp_classification_head.yaml index fadf8b7..c1e51db 100644 --- a/config/model_head/mlp_classification_head.yaml +++ b/config/model_head/mlp_classification_head.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* model_head: diff --git a/config/model_head/mlp_regression_head.yaml b/config/model_head/mlp_regression_head.yaml index 63f3158..a6b5e20 100644 --- a/config/model_head/mlp_regression_head.yaml +++ b/config/model_head/mlp_regression_head.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* model_head: diff --git a/config/model_head/patch_reconstruction_head.yaml b/config/model_head/patch_reconstruction_head.yaml index c07fd43..0bb1f5b 100644 --- a/config/model_head/patch_reconstruction_head.yaml +++ b/config/model_head/patch_reconstruction_head.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* model_head: diff --git a/config/model_head/sequence_classification_head.yaml b/config/model_head/sequence_classification_head.yaml index d793ca4..a1b5222 100644 --- a/config/model_head/sequence_classification_head.yaml +++ b/config/model_head/sequence_classification_head.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* model_head: diff --git a/config/task/finetune_regression_task_SCEReBrO.yaml b/config/task/finetune_regression_task_SCEReBrO.yaml index 567920a..4fd24a7 100644 --- a/config/task/finetune_regression_task_SCEReBrO.yaml +++ b/config/task/finetune_regression_task_SCEReBrO.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* task: diff --git a/config/task/finetune_task_SCEReBrO.yaml b/config/task/finetune_task_SCEReBrO.yaml index fbb7acc..76cb026 100644 --- a/config/task/finetune_task_SCEReBrO.yaml +++ b/config/task/finetune_task_SCEReBrO.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* task: diff --git a/config/task/pretrain_task_SCEReBrO.yaml b/config/task/pretrain_task_SCEReBrO.yaml index ad69397..1be7d03 100644 --- a/config/task/pretrain_task_SCEReBrO.yaml +++ b/config/task/pretrain_task_SCEReBrO.yaml @@ -15,7 +15,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #*----------------------------------------------------------------------------* task: diff --git a/criterion/ce_criterion.py b/criterion/ce_criterion.py index eed65a2..6b996f7 100644 --- a/criterion/ce_criterion.py +++ b/criterion/ce_criterion.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/criterion/focal_criterion.py b/criterion/focal_criterion.py index ee2e8c2..3605c20 100644 --- a/criterion/focal_criterion.py +++ b/criterion/focal_criterion.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/criterion/masked_reconstruction_loss.py b/criterion/masked_reconstruction_loss.py index 69f82df..af85a1e 100644 --- a/criterion/masked_reconstruction_loss.py +++ b/criterion/masked_reconstruction_loss.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/criterion/mse_criterion.py b/criterion/mse_criterion.py index 2aea9fa..9cfe9ca 100644 --- a/criterion/mse_criterion.py +++ b/criterion/mse_criterion.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/data_module/finetuning_data_module.py b/data_module/finetuning_data_module.py index e2de3b2..2b3ffe0 100644 --- a/data_module/finetuning_data_module.py +++ b/data_module/finetuning_data_module.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/data_module/pretraining_data_module.py b/data_module/pretraining_data_module.py index 40e6ec7..5ede1b2 100644 --- a/data_module/pretraining_data_module.py +++ b/data_module/pretraining_data_module.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/datasets/lmdb_dataset.py b/datasets/lmdb_dataset.py index ada3b04..7931f69 100644 --- a/datasets/lmdb_dataset.py +++ b/datasets/lmdb_dataset.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/datasets/tueg_dataset.py b/datasets/tueg_dataset.py index 146070c..138c86f 100644 --- a/datasets/tueg_dataset.py +++ b/datasets/tueg_dataset.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/common.py b/make_datasets/common.py index da3f8c3..c215c64 100644 --- a/make_datasets/common.py +++ b/make_datasets/common.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/electrode_positions.py b/make_datasets/electrode_positions.py index 5be2ec4..dde4ba8 100644 --- a/make_datasets/electrode_positions.py +++ b/make_datasets/electrode_positions.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_bci_ner.py b/make_datasets/make_bci_ner.py index 21f70c4..2cab2d5 100644 --- a/make_datasets/make_bci_ner.py +++ b/make_datasets/make_bci_ner.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_boas.py b/make_datasets/make_boas.py index 288bd94..ae21e56 100644 --- a/make_datasets/make_boas.py +++ b/make_datasets/make_boas.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_chb_mit.py b/make_datasets/make_chb_mit.py index 4dd8c7f..a2dda6f 100644 --- a/make_datasets/make_chb_mit.py +++ b/make_datasets/make_chb_mit.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_gwd.py b/make_datasets/make_gwd.py index f3b6999..7b17892 100644 --- a/make_datasets/make_gwd.py +++ b/make_datasets/make_gwd.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_isruc.py b/make_datasets/make_isruc.py index 8fac51c..d00be2e 100644 --- a/make_datasets/make_isruc.py +++ b/make_datasets/make_isruc.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_mental_arithmetic.py b/make_datasets/make_mental_arithmetic.py index d1f5f3a..5d91154 100644 --- a/make_datasets/make_mental_arithmetic.py +++ b/make_datasets/make_mental_arithmetic.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_mumtaz.py b/make_datasets/make_mumtaz.py index 44e587e..9c2ada9 100644 --- a/make_datasets/make_mumtaz.py +++ b/make_datasets/make_mumtaz.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_neonate.py b/make_datasets/make_neonate.py index 64014d8..54c219b 100644 --- a/make_datasets/make_neonate.py +++ b/make_datasets/make_neonate.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_physionet_mi.py b/make_datasets/make_physionet_mi.py index 69b6b08..88a75c6 100644 --- a/make_datasets/make_physionet_mi.py +++ b/make_datasets/make_physionet_mi.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_seed.py b/make_datasets/make_seed.py index bbb73e2..809740d 100644 --- a/make_datasets/make_seed.py +++ b/make_datasets/make_seed.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_seed_fra_ger.py b/make_datasets/make_seed_fra_ger.py index 4c146dc..a0cb4cf 100644 --- a/make_datasets/make_seed_fra_ger.py +++ b/make_datasets/make_seed_fra_ger.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_seed_v.py b/make_datasets/make_seed_v.py index 0a23450..f51f6cc 100644 --- a/make_datasets/make_seed_v.py +++ b/make_datasets/make_seed_v.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_seed_vig.py b/make_datasets/make_seed_vig.py index 7a5abf7..0de5b64 100644 --- a/make_datasets/make_seed_vig.py +++ b/make_datasets/make_seed_vig.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_shu_mi.py b/make_datasets/make_shu_mi.py index cad3319..d4346df 100644 --- a/make_datasets/make_shu_mi.py +++ b/make_datasets/make_shu_mi.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_sleepedfx.py b/make_datasets/make_sleepedfx.py index 07f2a2a..91ff8f9 100644 --- a/make_datasets/make_sleepedfx.py +++ b/make_datasets/make_sleepedfx.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_stew.py b/make_datasets/make_stew.py index 91df147..fb0c013 100644 --- a/make_datasets/make_stew.py +++ b/make_datasets/make_stew.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_tuab.py b/make_datasets/make_tuab.py index f783240..afe62ec 100644 --- a/make_datasets/make_tuab.py +++ b/make_datasets/make_tuab.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/make_datasets/make_tueg.py b/make_datasets/make_tueg.py index ab6de28..4c01543 100644 --- a/make_datasets/make_tueg.py +++ b/make_datasets/make_tueg.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/model_heads/__init__.py b/models/model_heads/__init__.py index 1f9c938..3cc5778 100644 --- a/models/model_heads/__init__.py +++ b/models/model_heads/__init__.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* 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 index 5b19f6a..4c4eaa0 100644 --- a/models/model_heads/mlp_classification_head.py +++ b/models/model_heads/mlp_classification_head.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/model_heads/mlp_regression_head.py b/models/model_heads/mlp_regression_head.py index fc983ac..0bfc211 100644 --- a/models/model_heads/mlp_regression_head.py +++ b/models/model_heads/mlp_regression_head.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/model_heads/patch_reconstruction_head.py b/models/model_heads/patch_reconstruction_head.py index efa7bbf..d654852 100644 --- a/models/model_heads/patch_reconstruction_head.py +++ b/models/model_heads/patch_reconstruction_head.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/model_heads/sequence_classification_head.py b/models/model_heads/sequence_classification_head.py index d61eb8f..56ad3eb 100644 --- a/models/model_heads/sequence_classification_head.py +++ b/models/model_heads/sequence_classification_head.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/modules/attention.py b/models/modules/attention.py index a5eebff..0b9e21d 100644 --- a/models/modules/attention.py +++ b/models/modules/attention.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/modules/patching.py b/models/modules/patching.py index 3c0eacc..207d544 100644 --- a/models/modules/patching.py +++ b/models/modules/patching.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/modules/pos_embed.py b/models/modules/pos_embed.py index b1aa4ac..9540dec 100644 --- a/models/modules/pos_embed.py +++ b/models/modules/pos_embed.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/models/s_cerebro.py b/models/s_cerebro.py index 85e33aa..9702f0a 100644 --- a/models/s_cerebro.py +++ b/models/s_cerebro.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/tasks/classification_task.py b/tasks/classification_task.py index fb15ef8..0c60c5d 100644 --- a/tasks/classification_task.py +++ b/tasks/classification_task.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/tasks/mae_pretraining.py b/tasks/mae_pretraining.py index 469169c..68fe6dd 100644 --- a/tasks/mae_pretraining.py +++ b/tasks/mae_pretraining.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/tasks/regression_task.py b/tasks/regression_task.py index 5fa0ba0..341c8ce 100644 --- a/tasks/regression_task.py +++ b/tasks/regression_task.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/tests/model_tests/test_attention.py b/tests/model_tests/test_attention.py index 6d7281e..26dff3e 100644 --- a/tests/model_tests/test_attention.py +++ b/tests/model_tests/test_attention.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * diff --git a/tests/model_tests/test_pipeline.py b/tests/model_tests/test_pipeline.py index 8f22f93..82ed761 100644 --- a/tests/model_tests/test_pipeline.py +++ b/tests/model_tests/test_pipeline.py @@ -14,7 +14,6 @@ #* See the License for the specific language governing permissions and * #* limitations under the License. * #* * -#* Author: Glenn Anta Bucagu * #* Author: BioFoundation Contributors * #* * #* Imported from the S-CEReBrO reference implementation (TimeFM). * From 333441cbe585d5feea6f2d2134ef15191577b4bb Mon Sep 17 00:00:00 2001 From: Glenn Bucagu <52401614+g12bftd@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:20:10 +0200 Subject: [PATCH 14/14] Delete make_datasets/make_dummy_scerebro_dataset.py --- make_datasets/make_dummy_scerebro_dataset.py | 262 ------------------- 1 file changed, 262 deletions(-) delete mode 100644 make_datasets/make_dummy_scerebro_dataset.py diff --git a/make_datasets/make_dummy_scerebro_dataset.py b/make_datasets/make_dummy_scerebro_dataset.py deleted file mode 100644 index ff4cedd..0000000 --- a/make_datasets/make_dummy_scerebro_dataset.py +++ /dev/null @@ -1,262 +0,0 @@ -#*----------------------------------------------------------------------------* -#* 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 * -#*----------------------------------------------------------------------------* - -"""Generate synthetic S-CEReBrO datasets for smoke-testing the training pipeline. - -This writes data in exactly the on-disk formats the real readers expect, so a run -started against it exercises the same code path as a run against prepared data. The -signals are band-limited noise and the labels are random: the output is for checking -that a pipeline runs, not for measuring anything. - -Two formats are produced. - -``LMDBDataset`` corpora - An LMDB directory per split whose values are pickled dictionaries with ``eeg``, - ``channel_coords`` and optionally ``label``, plus a ``{split}.meta.json`` index - listing the keys so start-up does not scan the database. - -``TUEGDataset`` corpus - A single LMDB whose values are fixed-size raw blobs holding a padded - ``float32[max_channels, timesteps]`` waveform followed by - ``float32[max_channels, 2, 3]`` coordinates, plus a newline-separated key list. - Padded channels are written with all-zero coordinates, which is how the reader - recovers ``num_padded_channels``. - -Example: - python -m make_datasets.make_dummy_scerebro_dataset --output $DATA_PATH -""" - -import argparse -import json -import os -import pickle -from pathlib import Path -from typing import Optional, Sequence - -import lmdb -import numpy as np - -SAMPLING_RATE = 200 -MAP_SIZE = 1 << 30 - -# Pre-training corpora named by config/data_module/pretrain_data_module_SCEReBrO.yaml. -PRETRAIN_CORPORA = ("SEED", "SEED-IV", "SEED-GER", "SEED-FRA", "BOAS", "GWD", "SleepEDFx", "BCI-NER") - - -def make_channel_coords(num_channels: int, rng: np.random.Generator) -> np.ndarray: - """Return plausible unit-sphere electrode-pair coordinates of shape (C, 2, 3). - - Each channel is described by two electrodes. The first is placed on a sphere of - unit radius and the second is a nearby point, which is what a bipolar derivation - or a scalp-electrode-plus-reference channel looks like to the encoder. - """ - theta = rng.uniform(0, np.pi, size=num_channels) - phi = rng.uniform(0, 2 * np.pi, size=num_channels) - first = np.stack( - [np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)], axis=-1 - ) - second = first + rng.normal(scale=0.15, size=(num_channels, 3)) - second /= np.linalg.norm(second, axis=-1, keepdims=True) - return np.stack([first, second], axis=1).astype(np.float32) - - -def make_eeg(num_channels: int, num_timesteps: int, rng: np.random.Generator) -> np.ndarray: - """Return band-limited noise of shape (C, T), loosely resembling scalp EEG.""" - t = np.arange(num_timesteps) / SAMPLING_RATE - signal = np.zeros((num_channels, num_timesteps), dtype=np.float32) - for freq in (2.0, 6.0, 10.0, 21.0): - amplitude = rng.uniform(0.2, 1.0, size=(num_channels, 1)) - phase = rng.uniform(0, 2 * np.pi, size=(num_channels, 1)) - signal += (amplitude * np.sin(2 * np.pi * freq * t + phase)).astype(np.float32) - signal += rng.normal(scale=0.3, size=signal.shape).astype(np.float32) - return signal * 20.0 - - -def write_lmdb_split( - directory: Path, - split: str, - num_samples: int, - num_channels: int, - num_timesteps: int, - rng: np.random.Generator, - label_kind: str = "none", - num_classes: int = 2, - sequence_length: Optional[int] = None, -) -> None: - """Write one ``{split}.lmdb`` of pickled samples plus its ``{split}.meta.json`` index. - - ``label_kind`` is one of ``none``, ``classification`` or ``regression``. It must be - ``none`` for pre-training corpora: those samples are concatenated with TUEG, which - carries no label, and a batch mixing samples with and without a ``label`` key - cannot be collated. - """ - if label_kind not in {"none", "classification", "regression"}: - raise ValueError(f"label_kind must be none, classification or regression, got {label_kind!r}") - directory.mkdir(parents=True, exist_ok=True) - path = directory / f"{split}.lmdb" - env = lmdb.open(str(path), map_size=MAP_SIZE) - keys = [] - - with env.begin(write=True) as txn: - for index in range(num_samples): - coords = make_channel_coords(num_channels, rng) - if sequence_length is None: - eeg = make_eeg(num_channels, num_timesteps, rng) - if label_kind == "classification": - label = int(rng.integers(num_classes)) - elif label_kind == "regression": - label = float(rng.uniform(0.0, 1.0)) - else: - label = None - else: - eeg = np.stack( - [make_eeg(num_channels, num_timesteps, rng) for _ in range(sequence_length)] - ) - label = ( - rng.integers(num_classes, size=sequence_length).astype(np.int64).tolist() - if label_kind == "classification" - else None - ) - - entry = {"eeg": eeg, "channel_coords": coords, "subject_id": f"S{index % 4:02d}"} - if label is not None: - entry["label"] = label - - key = f"{split}_{index:06d}".encode() - txn.put(key, pickle.dumps(entry)) - keys.append(key.decode()) - - env.close() - (directory / f"{split}.meta.json").write_text(json.dumps({"keys": keys})) - print(f" {path} ({num_samples} samples)") - - -def write_tueg( - directory: Path, - num_samples: int, - max_channels: int, - real_channels: int, - num_timesteps: int, - rng: np.random.Generator, -) -> None: - """Write the fixed-size-blob LMDB and key list that ``TUEGDataset`` reads.""" - directory.mkdir(parents=True, exist_ok=True) - path = directory / "TUEG.lmdb" - env = lmdb.open(str(path), map_size=MAP_SIZE) - keys = [] - - with env.begin(write=True) as txn: - for index in range(num_samples): - waveform = np.zeros((max_channels, num_timesteps), dtype=np.float32) - coords = np.zeros((max_channels, 2, 3), dtype=np.float32) - - # Vary how many channels are real so padding is genuinely exercised. - active = real_channels if index % 2 == 0 else max(4, real_channels // 2) - signal = make_eeg(active, num_timesteps, rng) - # TUEGDataset returns the stored waveform unchanged, so windows are - # min-max normalised offline exactly as make_tueg does. Without this the - # TUEG windows would be on a different scale from the LMDB corpora they - # are concatenated with. - minimum = signal.min(axis=1, keepdims=True) - maximum = signal.max(axis=1, keepdims=True) - waveform[:active] = ((signal - minimum) / (maximum - minimum + 1e-10) - 0.5) * 2.0 - coords[:active] = make_channel_coords(active, rng) - - key = f"tueg_{index:06d}".encode() - txn.put(key, waveform.tobytes() + coords.tobytes()) - keys.append(key.decode()) - - env.close() - (directory / "TUEG_keys.txt").write_text("\n".join(keys) + "\n") - print(f" {path} ({num_samples} windows, padded to {max_channels} channels)") - - -def main(argv: Optional[Sequence[str]] = None) -> None: - parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) - parser.add_argument( - "--output", required=True, - help="Root directory; pretraining/ and finetuning/ are created beneath it. " - "Use the same path you export as DATA_PATH.", - ) - parser.add_argument("--pretrain-samples", type=int, default=64, help="Windows per pre-training corpus.") - parser.add_argument("--finetune-samples", type=int, default=64, help="Windows in the fine-tuning train split.") - parser.add_argument("--max-channels", type=int, default=64, help="Channel capacity of the pre-training corpora.") - parser.add_argument("--slice-duration", type=int, default=30, help="Pre-training window length in seconds.") - parser.add_argument("--seed", type=int, default=0) - parser.add_argument( - "--datasets", nargs="+", default=["pretrain", "tuab"], - choices=["pretrain", "tuab", "isruc", "seed-vig"], - help="Which corpora to generate. 'pretrain' builds TUEG plus the eight LMDB corpora.", - ) - args = parser.parse_args(argv) - - rng = np.random.default_rng(args.seed) - root = Path(args.output).expanduser().resolve() - pretrain_timesteps = SAMPLING_RATE * args.slice_duration - - if "pretrain" in args.datasets: - print(f"Pre-training corpora -> {root/'pretraining'}") - write_tueg( - root / "pretraining" / "TUEG", args.pretrain_samples, - args.max_channels, args.max_channels // 2, pretrain_timesteps, rng, - ) - for corpus in PRETRAIN_CORPORA: - # Fewer channels than capacity, so the loader's zero-padding path runs. - write_lmdb_split( - root / "pretraining" / corpus, "train", args.pretrain_samples, - num_channels=args.max_channels // 2, num_timesteps=pretrain_timesteps, - rng=rng, label_kind="none", - ) - - if "tuab" in args.datasets: - print(f"TUAB (22 channels, 10 s, 2 classes) -> {root/'finetuning'/'TUAB'}") - for split, count in (("train", args.finetune_samples), - ("val", max(8, args.finetune_samples // 4)), - ("test", max(8, args.finetune_samples // 4))): - write_lmdb_split( - root / "finetuning" / "TUAB", split, count, - num_channels=22, num_timesteps=10 * SAMPLING_RATE, - rng=rng, label_kind="classification", num_classes=2, - ) - - if "isruc" in args.datasets: - print(f"ISRUC (6 channels, 30 s x 20 epochs, 5 classes) -> {root/'finetuning'/'ISRUC'}") - for split, count in (("train", args.finetune_samples // 4), - ("val", 8), ("test", 8)): - write_lmdb_split( - root / "finetuning" / "ISRUC", split, max(4, count), - num_channels=6, num_timesteps=30 * SAMPLING_RATE, - rng=rng, label_kind="classification", num_classes=5, sequence_length=20, - ) - - if "seed-vig" in args.datasets: - print(f"SEED-VIG (17 channels, 8 s, continuous) -> {root/'finetuning'/'SEED-VIG'}") - for split, count in (("train", args.finetune_samples), - ("val", 8), ("test", 8)): - write_lmdb_split( - root / "finetuning" / "SEED-VIG", split, count, - num_channels=17, num_timesteps=8 * SAMPLING_RATE, - rng=rng, label_kind="regression", - ) - - print(f"\nDone. Point DATA_PATH at {root}") - - -if __name__ == "__main__": - main()