DOCSEARCHbyAlgolia
Submit Docs

torchinfo

repository·main·Indexed 25 days ago

https://github.com/tyleryep/torchinfo

A tool for providing detailed model summaries for PyTorch networks, based on torchsummary. It allows developers to visualize layer shapes, parameter counts, and computational complexity (Mult-Adds) similar to TensorFlow's model.summary(). The library provides the summary() function for model analysis, with customizable output columns, depth control, and support for multiple inputs and data types.

Tokens
3.5K
Snippets
4
Records
40
Agent score
80%

What's inside torchinfo

  1. Use summary() to visualize PyTorch models

    main

    The summary() function provides a detailed visualization of your PyTorch model, including layer names, input/output shapes, parameter counts, and Mult-Adds.

    Note for Jupyter/Colab users: summary(model, ...) must be the returned value of the cell to display. If it is not, wrap it in print().

    from torchinfo import summary

model = ConvNet()
batch_size = 16
# Pass input_size as a tuple including the batch dimension
summary(model, input_size=(batch_size, 1, 28, 28))

  3. Get Model Summary as a string

    main

    To capture the summary output as a string (for logging or custom display) instead of printing it directly, set verbose=0 and cast the returned ModelStatistics object to a string.

    from torchinfo import summary

model_stats = summary(your_model, (1, 3, 28, 28), verbose=0)
summary_str = str(model_stats)
# summary_str contains the string representation of the summary!

  4. Summarize models with depth control

    main

    Use the depth parameter to control how many levels of the model hierarchy are displayed in the summary. This is useful for complex models like ResNet to avoid overly long outputs.

    import torchvision
from torchinfo import summary

model = torchvision.models.resnet152()
# depth=3 limits the hierarchy display
summary(model, (1, 3, 224, 224), depth=3)

  5. Use the summary function for model analysis

    main
    The summary function is the primary API for generating detailed reports of PyTorch models, including layer shapes, parameter counts, and estimated memory usage. You can pass the model and an input shape (as a tuple) to generate a basic summary.

  6. Configure summary() API parameters

    main

    The summary() function signature and its arguments:

    def summary(
    model: nn.Module,
    input_size: INPUT_SIZE_TYPE | None = None,
    input_data: INPUT_DATA_TYPE | None = None,
    batch_dim: int | None = None,
    cache_forward_pass: bool | None = None,
    col_names: Iterable[str] | None = None,
    col_width: int = 25,
    depth: int = 3,
    device: torch.device | str | None = None,
    dtypes: list[torch.dtype] | None = None,
    mode: str = "same",
    row_settings: Iterable[str] | None = None,
    verbose: int | None = None,
    **kwargs: Any,
) -> ModelStatistics:

    Key Arguments:

    • model: The PyTorch module to summarize.
    • input_size: Shape of input data (including batch size).
    • input_data: Actual tensors for the forward pass.
    • batch_dim: Index of the batch dimension. If None, assumes the first dimension is the batch.
    • col_names: Columns to display (e.g., "num_params", "mult_adds").
    • depth: How many nested layers to show.
    • device: The device to use for the model and inputs.
    • dtypes: List of dtypes if using input_size with non-FloatTensors.
    • mode: Determines if model.train() or model.eval() is called ("train", "eval", or "same").
    • verbose: 0 (quiet), 1 (default, print summary), 2 (show weights/bias in detail).

  7. Configure summary output columns and settings

    main

    You can customize the summary output by providing specific arguments:

    • dtypes: A list of torch data types for the inputs.
    • verbose: Controls the level of detail (e.g., verbose=2).
    • col_width: Sets the width for columns.
    • col_names: A list of column names to include (e.g., ["kernel_size", "output_size", "num_params", "mult_adds"]).
    • row_settings: A list of settings for rows (e.g., ["var_names"]).
    summary(
    model,
    (1, 100),
    dtypes=[torch.long],
    verbose=2,
    col_width=16,
    col_names=["kernel_size", "output_size", "num_params", "mult_adds"],
    row_settings=["var_names"],
)

  8. Use summary() with input data instead of input size

    main
    Instead of providing input_size, you can pass actual tensors via input_data. This is useful if your model's forward method accepts multiple arguments or complex data types.

  9. Summarize models with multiple inputs and different data types

    main
    For models requiring multiple inputs, pass a list of input shapes to the input_size argument and a corresponding list of dtypes. Alternatively, you can pass the actual input tensors directly to the input_data argument, and torchinfo will automatically infer the data types.

  10. Manage forward pass cache

    main

    If cache_forward_pass=True is set, torchinfo caches the result of the forward pass to speed up subsequent formatting changes (like changing depth or columns) in Jupyter Notebooks.

    To clear this cache, use clear_cached_forward_pass().

  11. Configure summary() output columns and rows

    main

    You can customize the information displayed in the summary table using col_names and row_settings.

    Supported col_names:

    • input_size
    • output_size
    • num_params
    • params_percent
    • kernel_size
    • groups
    • mult_adds
    • trainable

    Supported row_settings:

    • ascii_only
    • depth
    • var_names
    • hide_recursive_layers

  12. Configure model summary output with FormattingOptions

    main

    The FormattingOptions class controls how the model summary table is rendered, including column visibility, depth, and visual style.

    Key configuration attributes include:

    • max_depth: Limits how deep the summary tree goes.
    • verbose: Controls the level of detail (uses Verbosity enum).
    • col_names: A tuple of ColumnSettings determining which columns (e.g., INPUT_SIZE, NUM_PARAMS, MULT_ADDS) are displayed.
    • col_width: The fixed width for each data column.
    • row_settings: A set of RowSettings to toggle features like ASCII_ONLY, VAR_NAMES (show variable names), DEPTH (show depth index), or HIDE_RECURSIVE_LAYERS.
    • params_count_units, params_size_units, macs_units: Control the units used for parameters and MACs (uses Units enum).