← Back to API reference

neuralbench.main.Experiment

pydantic model neuralbench.main.Experiment

Brain-modeling experiment with support for loading pretrained weights.

Fields:

• augmentation (neuraltrain.augmentations.augmentations.BandRotationConfig | None)

• brain_model_config (neuraltrain.models.base.BaseModelConfig)

• brain_model_name (str)

• brain_model_output_size (int | None)

• compute_class_weights (bool)

• csv_config (neuraltrain.utils.CsvLoggerConfig | None)

• data (neuralbench.data.Data)

• delete_checkpoints_on_exit (bool)

• downstream_model_wrapper (neuralbench.modules.DownstreamWrapper | None)

• dummy (dict[str, Any])

• eval_only (bool)

• infra (exca.task.TaskInfra)

• lightning_optimizer_config (neuraltrain.optimizers.base.LightningOptimizer)

• loss (neuraltrain.losses.base.BaseLoss)

• metrics (list[neuraltrain.metrics.base.BaseMetric])

• pretrained_weights_fname (str | None)

• save_test_predictions (bool)

• seed (int)

• target_scaler (neuraltrain.utils.StandardScaler | None)

• task_name (str)

• test_full_metrics (list[neuraltrain.metrics.base.BaseMetric])

• test_full_retrieval_metrics (list[neuraltrain.metrics.base.BaseMetric])

• trainer_config (neuralbench.utils.TrainerConfig)

• validate_before_training (bool)

• wandb_config (neuraltrain.utils.WandbLoggerConfig | None)

field task_name: str = ''

field data: Data [Required]

field target_scaler: StandardScaler | None = None

field compute_class_weights: bool = False

field brain_model_config: BaseModelConfig [Required]

field brain_model_output_size: int | None = None

field pretrained_weights_fname: str | None = None

field downstream_model_wrapper: DownstreamWrapper | None = None

field trainer_config: TrainerConfig [Required]

field loss: BaseLoss [Required]

field lightning_optimizer_config: LightningOptimizer [Required]

field augmentation: BandRotationConfig | None = None

field eval_only: bool = False

field metrics: list[BaseMetric] [Required]

field validate_before_training: bool = True

field test_full_metrics: list[BaseMetric] = []

field test_full_retrieval_metrics: list[BaseMetric] = []

field save_test_predictions: bool = False

field csv_config: CsvLoggerConfig | None = None

field wandb_config: WandbLoggerConfig | None = None

field seed: int = 0

field delete_checkpoints_on_exit: bool = True

field infra: TaskInfra = TaskInfra(version='1')

field dummy: dict[str, Any] = {}

field brain_model_name: str = ''

prepare_pl_module(train_loader: DataLoader, val_loader: DataLoader | None = None) → None

fit(trainer: Trainer, train_loader: DataLoader, valid_loader: DataLoader) → None

setup_wandb_logger(wandb_config: WandbLoggerConfig, savedir: str) → WandbLogger

Setup wandb logger and launch initialization.

setup_run()

Setup paths and wandb logger.

setup_trainer(is_test: bool = False) → Trainer

Create callbacks and setup Trainer.

property run: dict[str, Any]

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.

Type:

Execute the full experiment lifecycle

test_predictions() → dict[str, Any]

Raw per-window test predictions.

Requires save_test_predictions=True. Returns a dict with:

• "metadata": a 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 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.

← Back to API reference