Neural Network Models

Implements factories for the neural network architectures (flax.nn.module objects) used in the main module of this package.

grannules.utils.model.model_from_params(hyperparams, n_outputs)

Creates a neural network model (flax.nn.Module) using the hyperparameters in hyperparams.

Parameters:

• hyperparams (dict[str, int | float | str]) –

  A dictionary with the following entries:

  • ’num_layers’, int, the number of layers the model should have.

  • ’dropout_rate’, float, the fraction of neurons to disable at random on any training run.

  • ’use_dropout_rate’, bool, whether to use a dropout rate at all.

  • ’layer_[n]_size’, int, the amount of neurons in the nth layer. There should be one of these arguments for every n in [0, num_layers).

  • ’layer_[n]_type’, str ‘relu’, ‘sigmoid’, or ‘tanh’ the operation performed by neurons in this layer. There should be one of these entries for every n in [0, num_layers), like for ‘layer_[n]_size’.

• n_outputs (int) – The number of outputs this neural network should have.

Returns:

Instance of created model (a subclass of flax.nn.Module).

Return type:

flax.nn.Module

grannules.utils.model.model_from_trial(trial, n_outputs)

Does the same thing as model_from_params(), but calls the suggest_ methods on trial to gather the hyperparameters. suggest_categorical() for the bool and str arguments, and suggest_int() and suggest_float() for the int and float arguments respectively.