网络结构可视化——torchinfo

torchinfo是一个用于PyTorch模型信息打印的Python包。它提供了一种简单而快速的方法来打印PyTorch模型的参数数量、计算图和内存使用情况等有用的信息，从而帮助深度学习开发人员更好地理解和优化他们的模型。整个模型的总参数数量和总内存使用情况。每个层的名称、输入形状、输出形状、参数数量和内存使用情况。如果指定了要显示的层数，则仅显示指定层数的信息。如果指定了show_input和s

文章共3,626字 · 阅读需要大约13分钟

一键AI生成摘要，助你高效阅读

问答

世事苍凉

4733人浏览 · 2023-04-13 18:18:17

世事苍凉 · 2023-04-13 18:18:17 发布

torchinfo介绍

torchinfo是一个用于PyTorch模型信息打印的Python包。它提供了一种简单而快速的方法来打印PyTorch模型的参数数量、计算图和内存使用情况等有用的信息，从而帮助深度学习开发人员更好地理解和优化他们的模型。

torchinfo可以打印以下模型的信息：

整个模型的总参数数量和总内存使用情况。
每个层的名称、输入形状、输出形状、参数数量和内存使用情况。
如果指定了要显示的层数，则仅显示指定层数的信息。
如果指定了show_input和show_output，则还会打印出中间层张量的形状和摘要信息。
如果指定了input_data，则还会显示模型的输出的内存使用情况。
如果使用了自定义处理程序，则可以打印自定义的格式。

在打印每个层的信息时，Torchinfo将使用以下格式：

==========================================================================================
Layer (type:depth-idx)                      Output Shape              Param #     Memory Usage
==========================================================================================
├─Conv2d: 1-1                               [1, 64, 112, 112]         9,408       9,408
├─BatchNorm2d: 1-2                          [1, 64, 112, 112]         128         512
├─ReLU: 1-3                                 [1, 64, 112, 112]         0           0
├─MaxPool2d: 1-4                            [1, 64, 56, 56]           0           0

其中，Layer列包括层的名称和类型，depth-idx列指示层的深度和在该深度中的索引，Output Shape列指示该层的输出形状，Param #列指示该层的参数数量，Memory Usage列指示该层的内存使用情况。

torchinfo安装

pip install torchinfo

或：

conda install -c conda-forge torchinfo

torchinfo使用

基础用法

使用Torchinfo非常简单，只需要将PyTorch模型作为参数传递给torchinfo.summary()函数即可。
例如：

import torch
import torch.nn as nn
from torchinfo import summary

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        x = self.relu(x)
        return x

model = Model()
summary(model, input_size=(1, 3, 224, 224))

这将打印出模型的详细信息，包括模型的名称、架构、参数数量、计算图和内存使用情况：

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Model                                    [1, 64, 224, 224]         --
├─Conv2d: 1-1                            [1, 64, 224, 224]         1,792
├─ReLU: 1-2                              [1, 64, 224, 224]         --
==========================================================================================
Total params: 1,792
Trainable params: 1,792
Non-trainable params: 0
Total mult-adds (M): 89.92
==========================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 25.69
Params size (MB): 0.01
Estimated Total Size (MB): 26.30
==========================================================================================

summary( ) 参数说明

参数说明
	model	要摘要的PyTorch模型
input_size	输入数据的形状，默认值为`None`
input_data	传递给模型正向传递的参数，默认值为`None`
batch_dim	输入数据的批次维度，默认值为`None`
cache_forward_pass	是否缓存模型正向传递的结果，默认值为`False`
col_names	输出中要显示哪些列，默认为`("output_size", "num_params")`
col_width	每列的宽度，默认值为`25`
depth	要显示的嵌套层数，默认值为`3`。
device	模型和输入数据使用的torch设备，默认情况下，使用`torch.cuda.is_available()`的结果
dtypes	如果您使用了`input_size`，则`torchinfo`假定您的输入使用`FloatTensors`。如果您的模型使用不同的数据类型，请指定该数据类型，对于多个输入，请指定两个输入的大小，并在此处指定每个参数的类型，默认值为`None`
mode	要使用的模型模式（`train`或`eval`），默认值为`"eval"`
row_settings	指定要在一行中显示哪些特性，默认为`("depth",)`
verbose	打印输出的详细程度。默认值为`1，`如果使用Juypter Notebook或Google Colab，则默认值为`0`
**kwargs	传递给`model.forward`函数的其他参数

补充用法

除了上面提到的基本功能，Torchinfo还提供了以下高级功能：

1、可以设置要显示的层数：您可以使用depth参数控制要显示的层数。默认情况下，Torchinfo将显示所有层的信息，但您可以使用depth参数来指定要显示的层数，例如：

summary(model, input_size=(1, 3, 224, 224), depth=2)

这将只显示前两层的信息。

2、可以打印出中间层张量的摘要信息：除了模型的输入和输出形状外，Torchinfo还可以打印出中间层张量的形状和摘要信息。要打印中间层张量的信息，请将show_input和show_output参数设置为True，例如：

summary(model, input_size=(1, 3, 224, 224), show_input=True, show_output=True)

这将显示所有中间层张量的形状和摘要信息。

3、可以显示模型的输出的内存使用情况：除了模型的参数数量和内存使用情况外，Torchinfo还可以显示模型的输出的内存使用情况。要显示模型的输出的内存使用情况，请将input_data参数设置为True，例如：

summary(model, input_size=(1, 3, 224, 224), input_data=True)

这将显示模型的输出的内存使用情况。

4、可以自定义打印格式：Torchinfo允许您自定义打印格式。您可以使用register_custom_handler()函数注册一个自定义处理程序，然后使用print_results参数来指定要使用的处理程序。例如：

from torchinfo.handlers import register_custom_handler

def custom_handler(module, input, output):
    print(f"Custom handler: {module} - input: {input[0].shape} - output: {output.shape}")

register_custom_handler(custom_handler)

summary(model, input_size=(1, 3, 224, 224), print_results=True)

这将使用自定义处理程序来打印模型信息。

5、支持多设备：如果您的模型在多个设备上运行，您可以使用device参数指定设备，以获得更准确的内存使用情况。例如：

summary(model, input_size=(1, 3, 224, 224), device='cuda')

这将使用CUDA设备计算内存使用情况。

6、支持多个输入：如果您的模型有多个输入，您可以使用input_size参数指定每个输入的形状。例如：

summary(model, input_size=[(1, 3, 224, 224), (1, 3, 224, 224)])

这将打印出模型的详细信息，并且假设模型有两个输入，每个输入形状为(1, 3, 224, 224)。

7、支持批量大小的动态调整：如果您的模型支持不同批量大小的输入，您可以使用batch_size参数指定批量大小，或者使用dynamic_batch=True来启用动态批量大小。例如：

summary(model, input_size=(1, 3, 224, 224), batch_size=16)

这将打印出模型的详细信息，并且假设批量大小为16。

8、支持忽略某些层：如果您想忽略模型中的某些层，您可以使用ignore_list参数来指定。例如：

summary(model, input_size=(1, 3, 224, 224), ignore_list=['conv'])

这将打印出模型的详细信息，但是不会包括名为'conv'的层的信息。

9、打印指定范围的层的信息：可以使用max_depth参数来指定要显示的层数。例如，为了打印第二到第六层的信息，可以将max_depth设置为6，并将depth_limit设置为2，以便从第二层开始打印。

summary(model, input_size=(1, 3, 224, 224), max_depth=6, depth_limit=2)

官方文档

https://github.com/TylerYep/torchinfohttps://github.com/TylerYep/torchinfo

def summary(
    model: nn.Module,
    input_size: Optional[INPUT_SIZE_TYPE] = None,
    input_data: Optional[INPUT_DATA_TYPE] = None,
    batch_dim: Optional[int] = None,
    cache_forward_pass: Optional[bool] = None,
    col_names: Optional[Iterable[str]] = None,
    col_width: int = 25,
    depth: int = 3,
    device: Optional[torch.device] = None,
    dtypes: Optional[List[torch.dtype]] = None,
    mode: str | None = None,
    row_settings: Optional[Iterable[str]] = None,
    verbose: int = 1,
    **kwargs: Any,
) -> ModelStatistics:
"""
Summarize the given PyTorch model. Summarized information includes:
    1) Layer names,
    2) input/output shapes,
    3) kernel shape,
    4) # of parameters,
    5) # of operations (Mult-Adds),
    6) whether layer is trainable

NOTE: If neither input_data or input_size are provided, no forward pass through the
network is performed, and the provided model information is limited to layer names.

Args:
    model (nn.Module):
            PyTorch model to summarize. The model should be fully in either train()
            or eval() mode. If layers are not all in the same mode, running summary
            may have side effects on batchnorm or dropout statistics. If you
            encounter an issue with this, please open a GitHub issue.

    input_size (Sequence of Sizes):
            Shape of input data as a List/Tuple/torch.Size
            (dtypes must match model input, default is FloatTensors).
            You should include batch size in the tuple.
            Default: None

    input_data (Sequence of Tensors):
            Arguments for the model's forward pass (dtypes inferred).
            If the forward() function takes several parameters, pass in a list of
            args or a dict of kwargs (if your forward() function takes in a dict
            as its only argument, wrap it in a list).
            Default: None

    batch_dim (int):
            Batch_dimension of input data. If batch_dim is None, assume
            input_data / input_size contains the batch dimension, which is used
            in all calculations. Else, expand all tensors to contain the batch_dim.
            Specifying batch_dim can be an runtime optimization, since if batch_dim
            is specified, torchinfo uses a batch size of 1 for the forward pass.
            Default: None

    cache_forward_pass (bool):
            If True, cache the run of the forward() function using the model
            class name as the key. If the forward pass is an expensive operation,
            this can make it easier to modify the formatting of your model
            summary, e.g. changing the depth or enabled column types, especially
            in Jupyter Notebooks.
            WARNING: Modifying the model architecture or input data/input size when
            this feature is enabled does not invalidate the cache or re-run the
            forward pass, and can cause incorrect summaries as a result.
            Default: False

    col_names (Iterable[str]):
            Specify which columns to show in the output. Currently supported: (
                "input_size",
                "output_size",
                "num_params",
                "params_percent",
                "kernel_size",
                "mult_adds",
                "trainable",
            )
            Default: ("output_size", "num_params")
            If input_data / input_size are not provided, only "num_params" is used.

    col_width (int):
            Width of each column.
            Default: 25

    depth (int):
            Depth of nested layers to display (e.g. Sequentials).
            Nested layers below this depth will not be displayed in the summary.
            Default: 3

    device (torch.Device):
            Uses this torch device for model and input_data.
            If not specified, uses result of torch.cuda.is_available().
            Default: None

    dtypes (List[torch.dtype]):
            If you use input_size, torchinfo assumes your input uses FloatTensors.
            If your model use a different data type, specify that dtype.
            For multiple inputs, specify the size of both inputs, and
            also specify the types of each parameter here.
            Default: None

    mode (str)
            Either "train" or "eval", which determines whether we call
            model.train() or model.eval() before calling summary().
            Default: "eval".

    row_settings (Iterable[str]):
            Specify which features to show in a row. Currently supported: (
                "ascii_only",
                "depth",
                "var_names",
            )
            Default: ("depth",)

    verbose (int):
            0 (quiet): No output
            1 (default): Print model summary
            2 (verbose): Show weight and bias layers in full detail
            Default: 1
            If using a Juypter Notebook or Google Colab, the default is 0.

    **kwargs:
            Other arguments used in `model.forward` function. Passing *args is no
            longer supported.

Return:
    ModelStatistics object
            See torchinfo/model_statistics.py for more information.
"""