MLPProbe#

class ts3.probes.MLPProbe(num_trials=100, num_folds=5, num_concurrent=4, seed=42, gpu_per_trial=0.25, cpu_per_trial=1, grad_clip=1.0)[source]#

Bases: ts3.probes.base.Probe

An MLP whose hyperparameters are swept with Optuna before the final fit.

Parameters:

  • num_trials (int) – sweep budget.

  • num_folds (int) – subject-wise folds; each trial fits on one and scores on the rest.

  • num_concurrent (int) – trials in flight at once.

  • seed (int) – the search, every trial and the final fit. Not the pretraining seed a submission is filed under, and not trial completion order, which is unseeded.

  • gpu_per_trial (float) – fractional GPU per trial.

  • cpu_per_trial (float) – CPUs per trial.

  • grad_clip (float | None) – max grad norm, null to disable. Fixed across trials, not swept.

fit_predict(train_embs, train_md, eval_embs, spec)[source]#

Fit on the pretrain units, return eval probabilities and what the fit learned.

The probabilities have shape (units, classes); the dict is whatever the fit is worth reporting about itself, empty for a probe with nothing to say.

Probabilities, not logits, and that is load-bearing: the multi-unit readout means neighbouring units together (ibl_bwb_eval.multi_unit.multi_unit_prediction()), and that rule is defined on probabilities. Columns follow spec.label_names; train_md is the TS3 unit table, so a probe that needs more than the region label (subject id, for a grouped split) reads it from there.

Return type:

tuple[ndarray, dict[str, Any]]