Function parameter_count_table

mmengine/analysis/complexity_analysis.py:295–357

Format the parameter count of the model (and its submodules or parameters) Adopted from https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/parameter_count.py Args: model (nn.Module): the model to count parameters. max_depth (int): maximum depth to recurs

(model: nn.Module, max_depth: int = 3)

def parameter_count_table(model: nn.Module, max_depth: int = 3) -> str:
    """Format the parameter count of the model (and its submodules or
    parameters)

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

    Args:
        model (nn.Module): the model to count parameters.
        max_depth (int): maximum depth to recursively print submodules or
            parameters

    Returns:
        str: the table to be printed
    """
    count: typing.DefaultDict[str, int] = parameter_count(model)
    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
    param_shape: typing.Dict[str, typing.Tuple] = {
        k: tuple(v.shape)
        for k, v in model.named_parameters()
    }

    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
    rows: typing.List[typing.Tuple] = []

    def format_size(x: int) -> str:
        if x > 1e8:
            return f'{x / 1e9:.1f}G'
        if x > 1e5:
            return f'{x / 1e6:.1f}M'
        if x > 1e2:
            return f'{x / 1e3:.1f}K'
        return str(x)

    def fill(lvl: int, prefix: str) -> None:
        if lvl >= max_depth:
            return
        for name, v in count.items():
            if name.count('.') == lvl and name.startswith(prefix):
                indent = ' ' * (lvl + 1)
                if name in param_shape:
                    rows.append(
                        (indent + name, indent + str(param_shape[name])))
                else:
                    rows.append((indent + name, indent + format_size(v)))
                    fill(lvl + 1, name + '.')

    rows.append(('model', format_size(count.pop(''))))
    fill(0, '')

    table = Table(
        title=f'parameter count of {model.__class__.__name__}', box=box.ASCII2)
    table.add_column('name')
    table.add_column('#elements or shape')

    for row in rows:
        table.add_row(*row)