Training

This module provides a script for training and evaluating JKOnet* and other models for learning diffusion terms on population data.

Functions

• numpy_collate

  A custom collate function for PyTorch’s DataLoader to properly stack or nest NumPy arrays when using JAX.

• main

  The main function that orchestrates the training loop, evaluation, logging, and visualization. It reads configurations, initializes models and datasets, and executes the training and evaluation processes.

Command-Line arguments

The script accepts the following command-line arguments:

• –solver, -s (EnumMethod):

  Name of the solver (model) to use. Choices are defined in the EnumMethod class.

• –dataset, -d (str):

  Name of the dataset to train the model on. The dataset should be prepared and located in a directory matching this name.

• –eval (str):

  Option to test the fit on ‘train_data’ or ‘test_data’ (e.g., for debugging purposes). Default is ‘test_data’.

• –wandb (bool):

  If specified, activates Weights & Biases logging for experiment tracking.

• –debug (bool):

  If specified, runs the script in debug mode (disables JIT compilation in JAX for easier debugging).

• –seed (int):

  Seed for random number generation to ensure reproducibility.

• –epochs (int):

  Number of epochs to train the model. If not specified, the number of epochs is taken from the configuration file.

Usage example

To train a model using the jkonet-star-potential solver on a dataset named my_dataset with wandb logging:

python train.py --solver jkonet-star-potential --dataset my_dataset --wandb

train.main(args: Namespace) → None

Main function to train a model on a specified dataset and evaluate it.

Parameters:

args (argparse.Namespace (see module description))

Returns:

The function trains the model, evaluates it, and optionally logs the results and metrics.

Return type:

None

train.numpy_collate(batch: List[ndarray | Tuple | List]) → ndarray | List

Collates a batch of samples into a single array or nested list of arrays.

This function recursively processes a batch of samples, stacking NumPy arrays, and collating lists or tuples by grouping elements together. If the batch consists of NumPy arrays, they are stacked. If the batch contains tuples or lists, the function recursively applies the collation.

This collate function is taken from the JAX tutorial with PyTorch Data Loading.

Parameters:

batch (List[Union[np.ndarray, Tuple, List]]) – A batch of samples where each sample is either a NumPy array, a tuple, or a list. It depends on the data loader.

Returns:

The collated batch, either as a stacked NumPy array or as a nested structure of arrays.

Return type:

np.ndarray