Skip to content

trainers

lanfactory.trainers.DatasetTorch 

DatasetTorch(file_ids, batch_size=32, label_lower_bound=None, label_upper_bound=None, features_key='data', label_key='labels', out_framework='torch')

Bases: Dataset

Dataset class for TorchMLP training.

Arguments 
file_ids (list):
    List of paths to the data files.
batch_size (int):
    Batch size.
label_lower_bound (float):
    Lower bound for the labels.
label_upper_bound (float):
    Upper bound for the labels.
features_key (str):
    Key for the features in the data files.
label_key (str):
    Key for the labels in the data files.
out_framework (str):
    Output framework.

lanfactory.trainers.LoadTorchMLP 

LoadTorchMLP(model_file_path=None, network_config=None, input_dim=None)

Class to load TorchMLP models.

Arguments 
model_file_path (str):
    Path to the model file.
network_config (dict):
    Network configuration.
input_dim (int):
    Input dimension.

lanfactory.trainers.LoadTorchMLPInfer 

LoadTorchMLPInfer(model_file_path=None, network_config=None, input_dim=None)

Class to load TorchMLP models for inference. (This was originally useful directly for application in the HDDM toolbox).

Arguments 
model_file_path (str):
    Path to the model file.
network_config (dict):
    Network configuration.
input_dim (int):
    Input dimension.

predict_on_batch 

predict_on_batch(x=None)

Intended as function that computes trial wise log-likelihoods from a matrix input. To be used primarily through the HDDM toolbox.

Arguments 
x (numpy.ndarray(dtype=numpy.float32)):
    Matrix which will be passed through the network.
    LANs expect the matrix columns to follow a specific order.
    When used in HDDM, x will be passed as follows.
    The first few columns are trial wise model parameters
    (order specified in the model_config file under the 'params' key).
    The last two columns are filled with trial wise
    reaction times and choices.
    When not used via HDDM, no such restriction applies.
Output 
numpy.ndarray(dtype = numpy.float32):
    Output of the network. When called through HDDM,
    this is expected as trial-wise log likelihoods
    of a given generative model.

lanfactory.trainers.MLPJax

Bases: Module

JaxMLP class.

Arguments 
layer_sizes (Sequence[int]):
    Sequence of integers containing the sizes of the layers.
activations (Sequence[str]):
    Sequence of strings containing the activation functions.
train (bool):
    Whether the model should be set to training mode or not.
train_output_type (str):
    The output type of the model during training.

__call__ 

__call__(inputs)

Call function for the JaxMLP class. Performs forward pass through the network.

Arguments 
inputs (jax.numpy.ndarray):
    Input tensor.

Returns:

  • jax.numpy.ndarray: –

    Output tensor.

load_state_from_file 

load_state_from_file(seed=42, input_dim=6, file_path=None)

Loads the state dictionary from a file.

Arguments 
seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
file_path (str):
    Path to the file containing the state dictionary.

Returns:

  • flax.core.frozen_dict.FrozenDict: –

    The state dictionary.

make_forward_partial 

make_forward_partial(seed=42, input_dim=6, state=None, add_jitted=False)

Creates a partial function for the forward pass of the network.

Arguments 
seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
state (flax.core.frozen_dict.FrozenDict):
    The state dictionary (if not loaded from file).
add_jitted (bool):
    Whether the partial function should be jitted or not.

Returns:

  • Callable: –

    The partial function for the forward pass of the network.

setup 

setup()

Setup function for the JaxMLP class. Initializes the layers and activation functions.

lanfactory.trainers.MLPJaxFactory 

MLPJaxFactory(network_config={}, train=True)

Factory function to create a MLPJax object.

Arguments 
network_config (dict):
    Dictionary containing the network configuration.
train (bool):
    Whether the model should be trained or not.

Returns:

  • MLPJax class initialized with the correct network configuration. –

lanfactory.trainers.ModelTrainerJaxMLP 

ModelTrainerJaxMLP(train_config=None, model=None, train_dl=None, valid_dl=None, allow_abs_path_folder_generation=False, pin_memory=False, seed=None)
Arguments 
train_config (dict):
    Dictionary containing the training configuration.
model (MLPJax):
    The MLPJax model to be trained.
train_dl (torch.utils.data.DataLoader):
    The training data loader.
valid_dl (torch.utils.data.DataLoader):
    The validation data loader.
allow_abs_path_folder_generation (bool):
    Whether the folder for the output files should be created or not.
pin_memory (bool):
    Whether the data loader should pin memory or not.
seed (int):
    Seed for the random number generator.

Returns:

  • ModelTrainerJaxMLP: –

    The ModelTrainerJaxMLP object.

create_train_state 

create_train_state(rng)

Create initial train state

run_epoch 

run_epoch(state, train=True, verbose=1, epoch=0, max_epochs=0)

Run one epoch of training or validation

Arguments 
state (flax.core.frozen_dict.FrozenDict):
    The state dictionary.
train (bool):
    Whether the model should is in training mode or not.
verbose (int):
    The verbosity level.
epoch (int):
    The current epoch.
max_epochs (int):
    The maximum number of epochs.

Returns:

  • tuple (flax.core.frozen_dict.FrozenDict, float): –

    The state dictionary and the mean epoch loss.

train_and_evaluate 

train_and_evaluate(output_folder='data/', output_file_id='fileid', run_id='runid', wandb_on=True, wandb_project_id='projectid', save_history=True, save_model=True, save_config=True, save_all=True, save_data_details=True, verbose=1)

Train and evaluate JAXMLP model.

Arguments 
output_folder (str):
    Path to the output folder.
output_file_id (str):
    The file id.
run_id (str):
    The run id.
wandb_on (bool):
    Whether to use wandb or not.
wandb_project_id (str):
    Project id for wandb.
save_history (bool):
    Whether to save the training history or not.
save_model (bool):
    Whether to save the model or not.
save_config (bool):
    Whether to save the training configuration or not.
save_all (bool):
    Whether to save all files or not.
save_data_details (bool):
    Whether to save the data details or not.
verbose (int):
    The verbosity level.

Returns:

  • flax.core.frozen_dict.FrozenDict: –

    The final state dictionary (model state).

lanfactory.trainers.ModelTrainerTorchMLP 

ModelTrainerTorchMLP(train_config=None, model=None, train_dl=None, valid_dl=None, allow_abs_path_folder_generation=False, pin_memory=True, seed=None)

train_config (dict):
    Training configuration.
model (TorchMLP):
    TorchMLP model.
train_dl (DatasetTorch):
    Training dataloader.
valid_dl (DatasetTorch):
    Validation dataloader.
allow_abs_path_folder_generation (bool):
    Whether to allow absolute path folder generation.
pin_memory (bool):
    Whether to pin memory (dataloader). Can affect speed.
seed (int):
    Random seed.

train_and_evaluate 

train_and_evaluate(output_folder='data/', output_file_id='fileid', run_id='runid', wandb_on=True, wandb_project_id='projectid', save_history=True, save_model=True, save_config=True, save_onnx=True, save_all=True, save_data_details=True, verbose=1)

Train and evaluate the model.

Arguments 
output_folder (str):
    Output folder.
output_file_id (str):
    Output file ID.
run_id (str):
    Run ID.
wandb_on (bool):
    Whether to use wandb.
wandb_project_id (str):
    Wandb project ID.
save_history (bool):
    Whether to save the training history.
save_model (bool):
    Whether to save the model.
save_config (bool):
    Whether to save the training configuration.
save_onnx (bool):
    Whether to save the model to ONNX format.
save_all (bool):
    Whether to save all.
save_data_details (bool):
    Whether to save the data details.
verbose (int):
    Verbosity level.

lanfactory.trainers.TorchMLP 

TorchMLP(network_config=None, input_shape=10, network_type=None, **kwargs)

Bases: Module

TorchMLP class.

Arguments 
network_config (dict):
    Network configuration.
input_shape (int):
    Input shape.
network_type (str):
    Network type.

forward 

forward(x)

Forward pass through network.

Arguments 
x (torch.Tensor):
    Input tensor.

Returns:

  • torch.Tensor: –

    Output tensor.

lanfactory.trainers.jax_mlp

MLPJax

Bases: Module

JaxMLP class.

Arguments 
layer_sizes (Sequence[int]):
    Sequence of integers containing the sizes of the layers.
activations (Sequence[str]):
    Sequence of strings containing the activation functions.
train (bool):
    Whether the model should be set to training mode or not.
train_output_type (str):
    The output type of the model during training.

__call__ 

__call__(inputs)

Call function for the JaxMLP class. Performs forward pass through the network.

Arguments 
inputs (jax.numpy.ndarray):
    Input tensor.

Returns:

  • jax.numpy.ndarray: –

    Output tensor.

load_state_from_file 

load_state_from_file(seed=42, input_dim=6, file_path=None)

Loads the state dictionary from a file.

Arguments 
seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
file_path (str):
    Path to the file containing the state dictionary.

Returns:

  • flax.core.frozen_dict.FrozenDict: –

    The state dictionary.

make_forward_partial 

make_forward_partial(seed=42, input_dim=6, state=None, add_jitted=False)

Creates a partial function for the forward pass of the network.

Arguments 
seed (int):
    Seed for the random number generator.
input_dim (int):
    Dimension of the input tensor.
state (flax.core.frozen_dict.FrozenDict):
    The state dictionary (if not loaded from file).
add_jitted (bool):
    Whether the partial function should be jitted or not.

Returns:

  • Callable: –

    The partial function for the forward pass of the network.

setup 

setup()

Setup function for the JaxMLP class. Initializes the layers and activation functions.

MLPJaxFactory 

MLPJaxFactory(network_config={}, train=True)

Factory function to create a MLPJax object.

Arguments 
network_config (dict):
    Dictionary containing the network configuration.
train (bool):
    Whether the model should be trained or not.

Returns:

  • MLPJax class initialized with the correct network configuration. –

ModelTrainerJaxMLP 

ModelTrainerJaxMLP(train_config=None, model=None, train_dl=None, valid_dl=None, allow_abs_path_folder_generation=False, pin_memory=False, seed=None)
Arguments 
train_config (dict):
    Dictionary containing the training configuration.
model (MLPJax):
    The MLPJax model to be trained.
train_dl (torch.utils.data.DataLoader):
    The training data loader.
valid_dl (torch.utils.data.DataLoader):
    The validation data loader.
allow_abs_path_folder_generation (bool):
    Whether the folder for the output files should be created or not.
pin_memory (bool):
    Whether the data loader should pin memory or not.
seed (int):
    Seed for the random number generator.

Returns:

  • ModelTrainerJaxMLP: –

    The ModelTrainerJaxMLP object.

create_train_state 

create_train_state(rng)

Create initial train state

run_epoch 

run_epoch(state, train=True, verbose=1, epoch=0, max_epochs=0)

Run one epoch of training or validation

Arguments 
state (flax.core.frozen_dict.FrozenDict):
    The state dictionary.
train (bool):
    Whether the model should is in training mode or not.
verbose (int):
    The verbosity level.
epoch (int):
    The current epoch.
max_epochs (int):
    The maximum number of epochs.

Returns:

  • tuple (flax.core.frozen_dict.FrozenDict, float): –

    The state dictionary and the mean epoch loss.

train_and_evaluate 

train_and_evaluate(output_folder='data/', output_file_id='fileid', run_id='runid', wandb_on=True, wandb_project_id='projectid', save_history=True, save_model=True, save_config=True, save_all=True, save_data_details=True, verbose=1)

Train and evaluate JAXMLP model.

Arguments 
output_folder (str):
    Path to the output folder.
output_file_id (str):
    The file id.
run_id (str):
    The run id.
wandb_on (bool):
    Whether to use wandb or not.
wandb_project_id (str):
    Project id for wandb.
save_history (bool):
    Whether to save the training history or not.
save_model (bool):
    Whether to save the model or not.
save_config (bool):
    Whether to save the training configuration or not.
save_all (bool):
    Whether to save all files or not.
save_data_details (bool):
    Whether to save the data details or not.
verbose (int):
    The verbosity level.

Returns:

  • flax.core.frozen_dict.FrozenDict: –

    The final state dictionary (model state).

lanfactory.trainers.torch_mlp

DatasetTorch 

DatasetTorch(file_ids, batch_size=32, label_lower_bound=None, label_upper_bound=None, features_key='data', label_key='labels', out_framework='torch')

Bases: Dataset

Dataset class for TorchMLP training.

Arguments 
file_ids (list):
    List of paths to the data files.
batch_size (int):
    Batch size.
label_lower_bound (float):
    Lower bound for the labels.
label_upper_bound (float):
    Upper bound for the labels.
features_key (str):
    Key for the features in the data files.
label_key (str):
    Key for the labels in the data files.
out_framework (str):
    Output framework.

LoadTorchMLP 

LoadTorchMLP(model_file_path=None, network_config=None, input_dim=None)

Class to load TorchMLP models.

Arguments 
model_file_path (str):
    Path to the model file.
network_config (dict):
    Network configuration.
input_dim (int):
    Input dimension.

LoadTorchMLPInfer 

LoadTorchMLPInfer(model_file_path=None, network_config=None, input_dim=None)

Class to load TorchMLP models for inference. (This was originally useful directly for application in the HDDM toolbox).

Arguments 
model_file_path (str):
    Path to the model file.
network_config (dict):
    Network configuration.
input_dim (int):
    Input dimension.

predict_on_batch 

predict_on_batch(x=None)

Intended as function that computes trial wise log-likelihoods from a matrix input. To be used primarily through the HDDM toolbox.

Arguments 
x (numpy.ndarray(dtype=numpy.float32)):
    Matrix which will be passed through the network.
    LANs expect the matrix columns to follow a specific order.
    When used in HDDM, x will be passed as follows.
    The first few columns are trial wise model parameters
    (order specified in the model_config file under the 'params' key).
    The last two columns are filled with trial wise
    reaction times and choices.
    When not used via HDDM, no such restriction applies.
Output 
numpy.ndarray(dtype = numpy.float32):
    Output of the network. When called through HDDM,
    this is expected as trial-wise log likelihoods
    of a given generative model.

ModelTrainerTorchMLP 

ModelTrainerTorchMLP(train_config=None, model=None, train_dl=None, valid_dl=None, allow_abs_path_folder_generation=False, pin_memory=True, seed=None)

train_config (dict):
    Training configuration.
model (TorchMLP):
    TorchMLP model.
train_dl (DatasetTorch):
    Training dataloader.
valid_dl (DatasetTorch):
    Validation dataloader.
allow_abs_path_folder_generation (bool):
    Whether to allow absolute path folder generation.
pin_memory (bool):
    Whether to pin memory (dataloader). Can affect speed.
seed (int):
    Random seed.

train_and_evaluate 

train_and_evaluate(output_folder='data/', output_file_id='fileid', run_id='runid', wandb_on=True, wandb_project_id='projectid', save_history=True, save_model=True, save_config=True, save_onnx=True, save_all=True, save_data_details=True, verbose=1)

Train and evaluate the model.

Arguments 
output_folder (str):
    Output folder.
output_file_id (str):
    Output file ID.
run_id (str):
    Run ID.
wandb_on (bool):
    Whether to use wandb.
wandb_project_id (str):
    Wandb project ID.
save_history (bool):
    Whether to save the training history.
save_model (bool):
    Whether to save the model.
save_config (bool):
    Whether to save the training configuration.
save_onnx (bool):
    Whether to save the model to ONNX format.
save_all (bool):
    Whether to save all.
save_data_details (bool):
    Whether to save the data details.
verbose (int):
    Verbosity level.

TorchMLP 

TorchMLP(network_config=None, input_shape=10, network_type=None, **kwargs)

Bases: Module

TorchMLP class.

Arguments 
network_config (dict):
    Network configuration.
input_shape (int):
    Input shape.
network_type (str):
    Network type.

forward 

forward(x)

Forward pass through network.

Arguments 
x (torch.Tensor):
    Input tensor.

Returns:

  • torch.Tensor: –

    Output tensor.