← Back to API reference

neuralbench.main.Experiment

pydantic model neuralbench.main.Experiment

Brain-modeling experiment with support for loading pretrained weights.

Fields:

• 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)

• 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 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 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(folder=None, cluster=None, logs='{folder}/logs/{user}/%j', job_name=None, timeout_min=None, nodes=1, tasks_per_node=1, cpus_per_task=None, gpus_per_node=None, mem_gb=None, max_pickle_size_gb=None, slurm_constraint=None, slurm_partition=None, slurm_account=None, slurm_qos=None, slurm_use_srun=False, slurm_additional_parameters=None, slurm_setup=None, conda_env=None, workdir=None, permissions=511, version='1', mode='cached', keep_in_ram=False)

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

← Back to API reference