# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import logging
import math
import platform
import resource
import time
import typing as tp
from pathlib import Path
import lightning.pytorch as pl
import numpy as np
import torch
import yaml
from exca import TaskInfra
from exca.cachedict import CacheDict
from lightning.pytorch.callbacks import (
Callback,
EarlyStopping,
LearningRateMonitor,
ModelCheckpoint,
)
from lightning.pytorch.loggers import CSVLogger, WandbLogger
from lightning.pytorch.loggers.logger import DummyLogger, Logger
from pydantic import model_validator
from torch.utils.data import DataLoader
from tqdm import tqdm
import neuralset as ns
from neuraltrain.augmentations import BandRotationConfig
from neuraltrain.losses import BaseLoss
from neuraltrain.metrics import BaseMetric
from neuraltrain.models.base import BaseModelConfig
from neuraltrain.optimizers import LightningOptimizer
from neuraltrain.utils import (
BaseExperiment,
CsvLoggerConfig,
StandardScaler,
WandbLoggerConfig,
)
from .aggregator import ( # noqa: F401
BenchmarkAggregator as BenchmarkAggregator,
)
from .callbacks import (
PlotConfusionMatrix,
PlotRegressionScatter,
PlotRegressionVectors,
RecordingLevelEval,
ResetPerStream,
TestFullRetrievalMetrics,
WindowPredictionCollector,
)
from .data import Data as Data # noqa: F401
from .model_factory import build_brain_model
from .modules import DownstreamWrapper
from .pl_module import BrainModule
from .utils import TrainerConfig, compute_class_weights_from_dataset
LOGGER = logging.getLogger(__name__)
[docs]
class Experiment(BaseExperiment):
"""Brain-modeling experiment with support for loading pretrained weights."""
task_name: str = ""
# Data
data: Data
target_scaler: StandardScaler | None = None
compute_class_weights: bool = False
# Model
brain_model_config: BaseModelConfig
brain_model_output_size: int | None = (
None # For convenience, e.g. when making yaml configs
)
pretrained_weights_fname: str | None = None
downstream_model_wrapper: DownstreamWrapper | None = None
# Optim
trainer_config: TrainerConfig
# Caps each epoch to a fresh random draw of this many hours of training
# windows, whatever the batch size or window duration.
max_neuro_hours_per_epoch: float | None = None
loss: BaseLoss
lightning_optimizer_config: LightningOptimizer
augmentation: BandRotationConfig | None = None
# Evaluation
eval_only: bool = False
metrics: list[BaseMetric]
validate_before_training: bool = True
test_full_metrics: list[BaseMetric] = []
test_full_retrieval_metrics: list[BaseMetric] = []
# When True, raw per-window test predictions/targets are folded into the
# cached ``run`` result (see ``WindowPredictionCollector``).
save_test_predictions: bool = False
# Weights & Biases
csv_config: CsvLoggerConfig | None = None
wandb_config: WandbLoggerConfig | None = None
# Internal properties
_brain_module: pl.LightningModule | None = None
_trainer: pl.Trainer | None = None
_csv_logger: CSVLogger | None = None
_wandb_logger: WandbLogger | None = None
_n_total_params: int | None = None
_n_trainable_params: int | None = None
# Others
seed: int = 0
delete_checkpoints_on_exit: bool = True
infra: TaskInfra = TaskInfra(version="1")
dummy: dict[str, tp.Any] = {} # Useful to avoid overwriting experiments between grids
brain_model_name: str = ""
space_name: str | None = None # data representation, e.g. a spaces/ variant
@model_validator(mode="after")
def _populate_brain_model_name(self) -> "Experiment":
"""Auto-populate brain_model_name from brain_model_config class when unset."""
if not self.brain_model_name:
self.brain_model_name = type(self.brain_model_config).__name__
return self
@model_validator(mode="after")
def _validate_metrics_num_classes(self) -> "Experiment":
"""Guard against stale num_classes in metrics after cross-file YAML merging.
YAML anchors are resolved at parse time within a single file. When a
dataset override changes ``brain_model_output_size`` without redefining
``metrics``, the inherited metric dicts still contain the old
``num_classes`` value. This validator catches such mismatches early.
"""
if self.brain_model_output_size is None:
return self
for metric in [*self.metrics, *self.test_full_metrics]:
kwargs = getattr(metric, "kwargs", None)
if kwargs is None:
continue
for key in ("num_classes", "num_labels"):
value = kwargs.get(key)
if value is not None and value != self.brain_model_output_size:
raise ValueError(
f"Metric '{metric.log_name}' has {key}={value} but "
f"brain_model_output_size={self.brain_model_output_size}. "
f"Dataset override likely changed brain_model_output_size "
f"without redefining metrics."
)
return self
[docs]
def prepare_pl_module(
self,
train_loader: DataLoader,
val_loader: DataLoader | None = None,
) -> None:
# Seed before any model construction so wrapper/adapters/lazy init do
# not depend on RNG already consumed by data preparation or setup.
pl.seed_everything(self.seed, workers=True)
(
brain_model,
self._n_total_params,
self._n_trainable_params,
ch_names,
) = build_brain_model(
brain_model_config=self.brain_model_config,
downstream_model_wrapper=self.downstream_model_wrapper,
pretrained_weights_fname=self.pretrained_weights_fname,
train_loader=train_loader,
val_loader=val_loader,
wandb_logger=self._wandb_logger,
loss=self.loss,
)
if self.target_scaler is not None:
assert hasattr(train_loader.dataset, "extractors")
neuro_extractor = train_loader.dataset.extractors.pop("neuro")
for batch in tqdm(train_loader, "Fitting target scaler"):
self.target_scaler.partial_fit(batch.data["target"])
if self.target_scaler._n_samples_seen > 5e5:
break
train_loader.dataset.extractors["neuro"] = neuro_extractor
loss_kwargs: dict[str, tp.Any] = {}
if self.compute_class_weights:
loss_kwargs, _ = compute_class_weights_from_dataset(
train_loader.dataset, # type: ignore[arg-type]
task=(
"multiclass"
if self.loss.__class__.__name__ == "CrossEntropyLoss"
else "multilabel"
),
logger=LOGGER,
)
self._brain_module = BrainModule(
model=brain_model,
ch_names=ch_names,
target_scaler=self.target_scaler,
augmentation=None if self.augmentation is None else self.augmentation.build(),
loss=self.loss.build(**loss_kwargs),
lightning_optimizer_config=self.lightning_optimizer_config,
metrics={metric.log_name: metric.build() for metric in self.metrics},
test_full_metrics={
metric.log_name: metric.build() for metric in self.test_full_metrics
},
test_full_retrieval_metrics={
metric.log_name: metric.build()
for metric in self.test_full_retrieval_metrics
},
)
pl.seed_everything(self.seed, workers=True)
[docs]
def fit(
self,
trainer: pl.Trainer,
train_loader: DataLoader,
valid_loader: DataLoader,
) -> None:
msg = "Prepare the BrainModule first with self.prepare_pl_module()"
assert self._brain_module is not None, msg
if self.validate_before_training:
LOGGER.info("Validating once before starting training...")
trainer.validate(model=self._brain_module, dataloaders=[valid_loader])
# Train model
trainer.fit(
model=self._brain_module,
train_dataloaders=train_loader,
val_dataloaders=valid_loader,
)
[docs]
def setup_wandb_logger(
self,
wandb_config: WandbLoggerConfig,
savedir: str,
) -> WandbLogger:
"""Setup wandb logger and launch initialization."""
import wandb
wandb.login(host=wandb_config.host)
logger = wandb_config.build(
save_dir=savedir,
xp_config=self.model_dump(),
)
try:
logger.experiment.config["_dummy"] = None # To launch initialization
except TypeError:
pass # Crashes if called in a second process, e.g. with DDP
return logger
[docs]
def setup_run(self):
"""Setup paths and wandb logger."""
savedir = self.infra.uid_folder()
LOGGER.info(f"UID folder: {savedir}")
# Save full config as yaml
if not savedir.exists():
savedir.mkdir(parents=True, exist_ok=False)
with open(savedir / "config.yaml", "w") as outfile:
yaml.dump(self.model_dump(), outfile, indent=4, default_flow_style=False)
if self.wandb_config is not None:
self._wandb_logger = self.setup_wandb_logger(self.wandb_config, str(savedir))
if self.csv_config is not None:
self._csv_logger = self.csv_config.build(save_dir=savedir)
[docs]
def setup_trainer(
self, is_test: bool = False, limit_train_batches: int | None = None
) -> pl.Trainer:
"""Create callbacks and setup Trainer, overriding ``limit_train_batches`` if given."""
callbacks: list[Callback] = []
if self.data.stream_by is not None:
callbacks.append(ResetPerStream())
neuro = self.data.neuro
if (
is_test
and isinstance(neuro, ns.extractors.MneRaw)
and neuro.scaler is not None
):
LOGGER.warning(
"Streamed test with data.neuro.scaler=%s, fit on each whole "
"recording: windows see future signal, so scores are optimistic "
"for a streamed submission.",
neuro.scaler,
)
if "confusion_matrix" in [metric.log_name for metric in self.metrics]:
labels: list[str] | None = None
if isinstance(self.data.target, ns.extractors.LabelEncoder):
ind_to_label: dict[int, str] = {}
for lbl, idx in self.data.target._label_to_ind.items():
ind_to_label.setdefault(idx, lbl)
labels = [ind_to_label[i] for i in sorted(ind_to_label)]
callbacks.append(PlotConfusionMatrix(labels=labels))
if is_test:
if self.save_test_predictions:
uid_folder = self.infra.uid_folder()
if uid_folder is None:
raise RuntimeError(
"save_test_predictions=True requires an infra cache "
"folder; configure ``infra.folder``."
)
callbacks.append(
WindowPredictionCollector(
folder=uid_folder / self._TEST_PREDICTIONS_DIR
)
)
if self.test_full_metrics:
callbacks.append(RecordingLevelEval())
if self.test_full_retrieval_metrics:
callbacks.append(
TestFullRetrievalMetrics(
event_type=self.data.target.event_types, # type: ignore[arg-type]
retrieval_set_sizes=(None, 250),
logger=LOGGER,
eval_val=False,
)
)
# Add regression vector visualization for multi-dimensional outputs
if (
self.brain_model_output_size is not None
and self.brain_model_output_size > 1
):
loss_name = (
self.loss.name
if hasattr(self.loss, "name")
else self.loss.__class__.__name__
)
if loss_name in ["MSELoss", "ClipLoss", "DistributedClipLoss"]:
callbacks.append(PlotRegressionVectors(num_samples=10))
# Add scatter plot for 1D regression outputs
if (
self.brain_model_output_size is not None
and self.brain_model_output_size == 1
):
loss_name = (
self.loss.name
if hasattr(self.loss, "name")
else self.loss.__class__.__name__
)
if loss_name == "MSELoss":
callbacks.append(PlotRegressionScatter())
else:
callbacks.append(LearningRateMonitor(logging_interval="step"))
callbacks.append(
EarlyStopping(
monitor=self.trainer_config.monitor,
patience=self.trainer_config.patience,
mode=self.trainer_config.mode,
verbose=True,
)
)
callbacks.append(
ModelCheckpoint(
dirpath=self.infra.uid_folder(),
filename="best",
monitor=self.trainer_config.monitor,
save_last=False,
mode=self.trainer_config.mode,
save_weights_only=True,
save_on_train_epoch_end=None,
save_top_k=1,
enable_version_counter=False,
)
)
loggers: list[Logger] = []
if self._wandb_logger is not None:
loggers.append(self._wandb_logger)
if self._csv_logger is not None:
loggers.append(self._csv_logger)
if not loggers:
loggers.append(DummyLogger())
trainer_config = self.trainer_config
if limit_train_batches is not None:
trainer_config = trainer_config.model_copy(
update={"limit_train_batches": limit_train_batches}
)
return trainer_config.build(
logger=loggers,
callbacks=callbacks,
accelerator="cpu" if self.infra.gpus_per_node == 0 else "auto",
devices=1 if is_test else self.infra.gpus_per_node,
num_nodes=1,
)
def _train_batch_limit(self, train_loader: DataLoader) -> int | None:
"""Per-rank ``limit_train_batches`` holding ``max_neuro_hours_per_epoch``.
``trainer_config.limit_train_batches`` is returned unchanged when the
training set is within the budget, and kept whenever it is stricter
(e.g. debug mode).
"""
assert self.max_neuro_hours_per_epoch is not None
configured = self.trainer_config.limit_train_batches
dataset = train_loader.dataset
segments = getattr(dataset, "segments", None)
if segments is not None:
total_seconds = sum(segment.duration for segment in segments)
elif self.data.duration is not None:
total_seconds = len(dataset) * self.data.duration # type: ignore[arg-type]
else:
raise ValueError(
"max_neuro_hours_per_epoch needs a fixed data.duration or segments"
)
total_hours = total_seconds / 3600
if total_hours <= self.max_neuro_hours_per_epoch:
return configured
# the loader is not sharded yet: under DDP each rank gets 1 / n_ranks
n_ranks = max(self.infra.gpus_per_node or 1, 1)
limit = math.ceil(
len(train_loader) * self.max_neuro_hours_per_epoch / total_hours / n_ranks
)
LOGGER.info(
"Capping training epochs to %d batches per rank (%.1f of %.1f window-hours)",
limit,
self.max_neuro_hours_per_epoch,
total_hours,
)
return limit if configured is None else min(limit, configured)
def _test(
self,
loaders: dict[str, DataLoader],
best_model_path: str | None,
) -> dict[str, float | None]:
"""Run the test phase on rank 0 with a dedicated trainer."""
tester = self.setup_trainer(is_test=True)
return dict(
tester.test(
self._brain_module,
dataloaders=loaders["test"],
ckpt_path=best_model_path,
)[0]
)
def _cleanup(self, trainer: pl.Trainer) -> None:
"""Delete checkpoint and finalize W&B."""
if (
# Deleting from another rank can strip the checkpoint before rank
# zero has tested with it.
trainer.global_rank == 0
and self.delete_checkpoints_on_exit
and not self.eval_only
and hasattr(trainer.checkpoint_callback, "best_model_path")
):
best_model_path = getattr(
trainer.checkpoint_callback, "best_model_path", None
)
assert best_model_path is not None
Path(best_model_path).unlink(missing_ok=True)
LOGGER.info("Deleted checkpoint: %s", best_model_path)
if self._wandb_logger is not None:
import wandb
wandb.finish()
def _exclude_from_cache_uid(self) -> list[str]:
"""Config paths that must not affect the cache key.
``space_name`` only labels a run: the ``spaces/`` overlay it names
already changes the config. The LoRA fields are inert while
``lora_config`` is unset, so dropping them then keeps finetune / probe
runs on the cache keys they had before LoRA support landed.
"""
excluded = ["wandb_config", "csv_config", "brain_model_name", "space_name"]
wrapper = self.downstream_model_wrapper
if wrapper is not None and wrapper.lora_config is None:
excluded += [
"downstream_model_wrapper.lora_config",
"downstream_model_wrapper.lora_target_modules",
]
return excluded
@infra.apply(exclude_from_cache_uid="method:_exclude_from_cache_uid")
def run(self) -> dict[str, tp.Any]:
"""Execute the full experiment lifecycle: setup, train, test, cleanup.
Returns a dict of test metrics (e.g. ``{"test/bal_acc": 0.85, ...}``)
plus ``n_total_params`` and ``n_trainable_params``.
"""
logging.basicConfig(level=logging.INFO)
logging.getLogger("numexpr").setLevel(logging.WARNING)
logging.getLogger("fontTools").setLevel(logging.WARNING)
logging.getLogger("fontTools.subset").setLevel(logging.WARNING)
logging.getLogger("fontTools.ttLib").setLevel(logging.WARNING)
logging.getLogger("matplotlib.font_manager").setLevel(logging.WARNING)
logging.getLogger("exca").propagate = False
logging.getLogger("neuralset").propagate = False
# Defense-in-depth seeding at the top of ``run()``. Data-side
# determinism (shuffle, weighted sampler, worker RNGs) is driven by
# ``Data.seed`` via explicit ``torch.Generator``s and is immune to
# the global RNG state when set.
pl.seed_everything(self.seed, workers=True)
self.setup_run()
loaders = self.data.prepare()
limit_train_batches = None
if self.max_neuro_hours_per_epoch is not None:
limit_train_batches = self._train_batch_limit(loaders["train"])
trainer = self.setup_trainer(limit_train_batches=limit_train_batches)
self.prepare_pl_module(loaders["train"], loaders.get("val"))
test_results: dict[str, tp.Any] = {}
best_model_path: str | None = None
training_time_s: float | None = None
peak_gpu_memory_mb: float | None = None
peak_cpu_memory_mb: float | None = None
if not self.eval_only:
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
t0 = time.perf_counter()
self.fit(trainer, loaders["train"], loaders["val"])
training_time_s = time.perf_counter() - t0
if torch.cuda.is_available():
peak_gpu_memory_mb = torch.cuda.max_memory_allocated() / (1024**2)
rusage = resource.getrusage(resource.RUSAGE_SELF)
# macOS reports bytes, Linux reports KB
divisor = 1024 if platform.system() != "Darwin" else 1024**2
peak_cpu_memory_mb = rusage.ru_maxrss / divisor
# Only rank zero writes the checkpoint, and it deletes it in
# ``_cleanup``, so reading it elsewhere races that deletion.
if trainer.global_rank == 0 and isinstance(
trainer.checkpoint_callback, ModelCheckpoint
):
best_model_path = trainer.checkpoint_callback.best_model_path
assert best_model_path is not None
best_ckpt = torch.load(
best_model_path,
map_location=torch.device("cpu"),
weights_only=True,
)
best_epoch = best_ckpt["epoch"]
LOGGER.info("Best epoch: %i", best_epoch)
for logger in [self._wandb_logger, self._csv_logger]:
if logger is not None:
logger.log_metrics({"best_epoch": best_epoch})
if (
getattr(self.infra, "gpus_per_node", 0) > 1
and torch.distributed.is_initialized()
):
torch.distributed.destroy_process_group()
if trainer.global_rank == 0:
# ``WindowPredictionCollector`` (added in ``setup_trainer`` when
# ``save_test_predictions`` is set) streams the raw per-window
# predictions to ``infra.uid_folder()/<_TEST_PREDICTIONS_DIR>`` as a
# byproduct of the test loop. They live next to ``job.pkl`` so they
# survive cache hits and stay out of the cached result dict; read
# them back via ``test_predictions()``.
test_results.update(self._test(loaders, best_model_path))
self._cleanup(trainer)
test_results.update(
{
"n_total_params": self._n_total_params,
"n_trainable_params": self._n_trainable_params,
"training_time_s": training_time_s,
"peak_gpu_memory_mb": peak_gpu_memory_mb,
"peak_cpu_memory_mb": peak_cpu_memory_mb,
}
)
return test_results
# Subfolder (under ``infra.uid_folder()``) where ``WindowPredictionCollector``
# streams the per-window test prediction artifacts via exca's ``CacheDict``.
_TEST_PREDICTIONS_DIR: tp.ClassVar[str] = "test_predictions"
[docs]
def test_predictions(self) -> dict[str, tp.Any]:
"""Raw per-window test predictions.
Requires ``save_test_predictions=True``. Returns a dict with:
- ``"metadata"``: a :class:`pandas.DataFrame` with one row per test
window (``timeline``, ``batch_idx``, ``dataloader_idx``, plus
``subject_id`` and a retrieval ``group`` label when available);
- ``"y_true"`` / ``"y_pred"``: arrays of shape ``(n_windows, ...)``
aligned with ``metadata``, concatenated across batches.
``WindowPredictionCollector`` streams the predictions to the uid folder
during the test loop (metadata as CSV, arrays appended to a shared
memmap file), so they survive cache hits and are read straight from
disk. This is a read-only accessor: call :meth:`run` first (a cache hit
is fine); if the artifacts are missing it raises rather than launching a
run. The per-batch array chunks are concatenated on read, so loading
materializes the full arrays in RAM.
"""
uid_folder = self.infra.uid_folder()
if uid_folder is None:
raise RuntimeError(
"Cannot read test predictions without an infra cache folder; "
"configure ``infra.folder``."
)
cache: CacheDict = CacheDict(folder=uid_folder / self._TEST_PREDICTIONS_DIR)
keys = set(cache.keys())
if not keys:
raise ValueError(
"Test predictions are unavailable. Set "
"save_test_predictions=True and (re)run the experiment "
"(toggling the flag yields a fresh cache entry)."
)
def _concat(prefix: str) -> np.ndarray:
chunks = sorted(key for key in keys if key.startswith(prefix))
return np.concatenate([np.asarray(cache[key]) for key in chunks], axis=0)
collector = WindowPredictionCollector
return {
"metadata": cache[collector._METADATA_KEY],
"y_true": _concat(collector._Y_TRUE_PREFIX),
"y_pred": _concat(collector._Y_PRED_PREFIX),
}
# BenchmarkAggregator lives in aggregator.py and forward-references Experiment
# (under TYPE_CHECKING to avoid a circular import). Pydantic needs the real
# class at validation time, so we rebuild the schema here once both are defined.
BenchmarkAggregator.model_rebuild(_types_namespace={"Experiment": Experiment})