Skip to content

ezpz.models⚓︎

Model zoo for ezpz examples and benchmarks.

Submodules expose Llama-style transformers, vision transformers, attention primitives, and a minimal MLP used by smoke tests.

summarize_model(model, verbose=False, depth=1, input_size=None) ⚓︎

Print a torchinfo model summary if available.

Parameters:

Name Type Description Default
model Module

The nn.Module to summarise.

 required
verbose bool

Passed through to torchinfo.summary.

 False
depth int

Depth of nested module display.

 1
input_size Optional[Sequence[int]]

Optional input size for shape inference.

 None

Returns:

Type Description
object | None

The torchinfo.ModelStatistics object, or None if
object | None

torchinfo is not installed.
Source code in src/ezpz/models/__init__.py 
def summarize_model(
    model: "torch.nn.Module",
    verbose: bool = False,
    depth: int = 1,
    input_size: Optional[Sequence[int]] = None,
) -> object | None:
    """Print a ``torchinfo`` model summary if available.

    Args:
        model: The ``nn.Module`` to summarise.
        verbose: Passed through to ``torchinfo.summary``.
        depth: Depth of nested module display.
        input_size: Optional input size for shape inference.

    Returns:
        The ``torchinfo.ModelStatistics`` object, or ``None`` if
        ``torchinfo`` is not installed.
    """
    try:
        import torchinfo
        from torchinfo import summary

        summary_str = summary(
            model,
            input_size=input_size,
            depth=depth,
            verbose=verbose,
        )
        # logger.info(f'\n{summary_str}')
        return summary_str

    except (ImportError, ModuleNotFoundError):
        logger.warning(
            'torchinfo not installed, unable to print model summary!'
        )