NDT

class ts2.models.single_session.NDT(max_spikes, unit_emb_dim, encoder_num_heads, encoder_ffn_factor, encoder_num_layers, encoder_dropout, pre_encoder_dropout, post_encoder_dropout, encoder_activation, use_custom_init=False, initrange=0.1, tfixup_scale_base=None, tfixup_v_scale_factor=None)

Bases: core.model.BaseModel

Single-session Transformer encoder over binned spikes [Ye and Pandarinath, 2021].

Reference implementation: neural-data-transformers.

link_datasets() must be called first: the hidden dimension is num_units * max(unit_emb_dim, 1), so the encoder is built there.

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

Parameters:

• max_spikes (int) – Size of the count embedding table; counts are clipped to one below it.

• unit_emb_dim (int) – Per-unit count embedding width. 0 feeds raw counts, the NDT default.

• encoder_num_heads (int) – Number of attention heads.

• encoder_ffn_factor (int) – Feed-forward sublayer width multiplier relative to the hidden dim.

• encoder_num_layers (int) – Number of Transformer encoder layers.

• encoder_dropout (float) – Dropout inside each encoder layer.

• pre_encoder_dropout (float) – Dropout applied to embeddings before the encoder.

• post_encoder_dropout (float) – Dropout applied to encoder output.

• encoder_activation (str) – Activation of the feed-forward sublayer.

• use_custom_init (bool) – If True, apply T-Fixup weight initialisation (see custom_init()).

• initrange (float) – Uniform init range for non-Transformer weights.

• tfixup_scale_base (Optional[float]) – T-Fixup base scale factor (required when use_custom_init=True).

• tfixup_v_scale_factor (Optional[float]) – Additional scale applied to value weights (required when use_custom_init=True).

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, positions)

Run the model on one collated batch.

Takes whatever input_fn() put under model_inputs, splatted in as keyword arguments, and returns what the trainer’s loss reads.

custom_init()

T-Fixup init for the encoder; see core.nn.tfixup_init_().

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.