Function parameter_count

mmengine/analysis/complexity_analysis.py:270–292

Count parameters of a model and its submodules. Adopted from https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/parameter_count.py Args: model (nn.Module): the model to count parameters. Returns: dict[str, int]: the key is either a parameter name or a m

(model: nn.Module)

def parameter_count(model: nn.Module) -> typing.DefaultDict[str, int]:
    """Count parameters of a model and its submodules.

    Adopted from
    https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/parameter_count.py

    Args:
        model (nn.Module): the model to count parameters.

    Returns:
        dict[str, int]: the key is either a parameter name or a module name.
        The value is the number of elements in the parameter, or in all
        parameters of the module. The key "" corresponds to the total
        number of parameters of the model.
    """
    count = defaultdict(int)  # type: typing.DefaultDict[str, int]
    for name, param in model.named_parameters():
        size = param.numel()
        name = name.split('.')
        for k in range(0, len(name) + 1):
            prefix = '.'.join(name[:k])
            count[prefix] += size
    return count


def parameter_count_table(model: nn.Module, max_depth: int = 3) -> str: