count_parameters

selfeeg.utils.utils.count_parameters(model: Module, return_table: bool = False, print_table: bool = False, add_not_trainable=False) → [<class 'int'>, typing.Optional[pandas.core.frame.DataFrame]][source]

counts the number of trainable parameters of a Pytorch’s nn.Module.

It can additionally create a two column dataframe with module’s name and number of trainable parameters. Not trainable parameters can be also added to the table if specified.

The implementation is an enriched implementation inspired from [stacko1] and [stacko2] .

Parameters:

  • model (nn.Module) – The model to scroll.

  • return_table (bool, optional) –

    Whether to return a with module’s name and number of trainable parameters or not.

    Default = False

  • print_table (bool, optional) –

    Whether to print the created table or not.

    Default = False

  • add_not_trainable (bool, optional) –

    Whether to add blocks with 0 trainable parameters to the table or not.

    Default = False

Returns:

  • total_params (int) – The number of trainable parameters.

  • layer_table (pd.DataFrame, optional) – A two column dataframe with module’s name and number of trainable parameters.

References

Example

>>> import selfeeg.utils
>>> import selfeeg.models
>>> mdl = models.ShallowNet(4,8,1024)
>>> for n, i in enumerate(mdl.parameters()): # bias require grad put to False
...     i.requires_grad=False if n in [1,3,5,7] else True
>>> a,b = utils.count_parameters(mdl, True,True,True)
>>> print (b == 23760) # should return True