count_parameters

class brainstate.nn.count_parameters(module, precision=2, return_table=False)

Count and display the number of trainable parameters in a neural network model.

This function iterates through all the parameters of the given model, counts the number of parameters for each module, and displays them in a table. It also calculates and returns the total number of trainable parameters.

Parameters:

modelbrainstate.nn.Module

The neural network model for which to count parameters.

Returns:

int

The total number of trainable parameters in the model.

Prints:

A pretty-formatted table showing the number of parameters for each module, followed by the total number of trainable parameters.