AutoencoderMLP

class ts2.models.single_session.AutoencoderMLP(encoder_depth=2, decoder_depth=2, hidden_dim=256, dropout=0.2, activation='relu', batch_norm=True)

Bases: core.model.BaseModel

Simple Autoencoder with MLP encoder/decoder for TS2 neural prediction.

Takes (masked) binned spike counts of shape (B, T, N) as input and outputs log-rates of shape (B, T, N) for the Poisson NLL reconstruction loss.

N (number of neurons) is variable per recording; it is resolved lazily in link_datasets so that a single output layer can be created once the dataset is available.

Not tested at any context_length other than the default (1.0); see Variable context.

Parameters:

• encoder_depth (int) – Number of encoder layers. Each halves the width after the first.

• decoder_depth (int) – Number of decoder layers. The first keeps the bottleneck width, each later one doubles it.

• hidden_dim (int) – Width of the first encoder layer.

• dropout (float) – Dropout probability applied after each activation.

• activation (Literal['relu', 'gelu', 'tanh']) – Pointwise non-linearity.

• batch_norm (bool) – If True, insert BatchNorm1d after each linear layer.

link_datasets(train_dataset, val_dataset, test_dataset=None)

Build whatever the model sizes from the datasets.

Called before configure_readout(), so what is read here sizes the readout.

Parameters:

• train_dataset (IBLBrainWideBenchTS2) – Training split, the one to size from.

• val_dataset (IBLBrainWideBenchTS2) – Validation split.

• test_dataset (Optional[IBLBrainWideBenchTS2]) – Test split, or None when the run scores none.

input_fn(data)

Convert one trial’s data into the model’s inputs.

Runs per item on the dataloader workers, as a dataset transform, so it returns tensors rather than batches.

Parameters:

data (Data) – One trial, exposing the interval key configure_readout() set.

Return type:

dict

Returns:

Dict whose model_inputs entry is splatted into forward(). Extra top-level keys are the model’s own to read in its trainer.

forward(spikes)

Predict log spike rates from masked spike counts.

Parameters:

spikes (Tensor) – (B, T, N) masked binned spike counts

Returns:

(B, T, N) predicted log spike rates

Return type:

log_rates

classmethod create_search_space(trial, cfg)

Map out the model’s Optuna search space.

Call trial.suggest_*; the names suggested become the keys process_tunable_params() receives.

Parameters:

• trial (Trial) – Optuna trial to register suggestions on.

• cfg (DictConfig) – The run config, for values the space depends on.

classmethod process_tunable_params(tune_params)

Turn suggested hyperparameters into config overrides.

Runs before the config is filled, so this is where a suggestion is mapped onto the config path it sets (batch_size_log2 -> batch_size), a value is derived from another, or a default is supplied for something not being tuned.

Parameters:

tune_params (dict) – The names create_search_space() suggested, with their values.

Return type:

dict

Returns:

The overrides to apply to the config. The default returns them unchanged.