torchinfo

（原名 torch-summary）

Torchinfo 提供了与 PyTorch 中 print(your_model) 所提供信息互补的信息，类似于 TensorFlow 的 model.summary() API，用于查看模型的可视化结构，在调试网络时非常有帮助。在这个项目中，我们实现了类似的功能，并为 PyTorch 创建了一个简洁、易用的接口，方便在项目中使用。

这是由 @sksq96 和 @nmhkahn 原创的 torchsummary 和 torchsummaryX 项目的完全重写版本。通过引入全新的 API，本项目解决了原项目遗留的所有问题和拉取请求。

支持 PyTorch 1.4.0 及以上版本。

使用方法

pip install torchinfo

或者通过 conda 安装：

conda install -c conda-forge torchinfo

使用方式

from torchinfo import summary

model = ConvNet()
batch_size = 16
summary(model, input_size=(batch_size, 1, 28, 28))

================================================================================================================
Layer (type:depth-idx)          Input Shape          Output Shape         Param #            Mult-Adds
================================================================================================================
SingleInputNet                  [7, 1, 28, 28]       [7, 10]              --                 --
├─Conv2d: 1-1                   [7, 1, 28, 28]       [7, 10, 24, 24]      260                1,048,320
├─Conv2d: 1-2                   [7, 10, 12, 12]      [7, 20, 8, 8]        5,020              2,248,960
├─Dropout2d: 1-3                [7, 20, 8, 8]        [7, 20, 8, 8]        --                 --
├─Linear: 1-4                   [7, 320]             [7, 50]              16,050             112,350
├─Linear: 1-5                   [7, 50]              [7, 10]              510                3,570
================================================================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
Total mult-adds (M): 3.41
================================================================================================================
Input size (MB): 0.02
Forward/backward pass size (MB): 0.40
Params size (MB): 0.09
Estimated Total Size (MB): 0.51
================================================================================================================

注意：如果您使用的是 Jupyter Notebook 或 Google Colab，summary(model, ...) 必须是单元格的返回值。如果不是，您应该将 summary 包裹在 print() 中，例如 print(summary(model, ...))。示例请参见 tests/jupyter_test.ipynb。

此版本现在支持：

• RNN、LSTM 及其他递归层
• 分支输出，可用于按指定深度探索模型层
• 返回包含所有摘要数据字段的 ModelStatistics 对象
• 可配置的行数和列数
• Jupyter Notebook / Google Colab

其他新特性：

• 详细模式，可显示权重和偏置层
• 支持输入数据或仅输入形状！
• 可自定义线条宽度和批量维度
• 全面的单元测试、输出测试、代码风格检查和代码覆盖率测试

社区贡献：

• Sequentials 和 ModuleLists（感谢 @roym899）
• 改进的 Mult-Add 计算（感谢 @TE-StefanUhlich、@zmzhang2000）
• 字典/其他格式的输入数据（感谢 @e-dorigatti）
• 剪枝层支持（感谢 @MajorCarrot）

文档

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 = "same",
    row_settings: Optional[Iterable[str]] = None,
    verbose: int = 1,
    **kwargs: Any,
) -> ModelStatistics:
"""
总结给定的 PyTorch 模型。总结的信息包括：
    1) 层名称，
    2) 输入/输出形状，
    3) 卷积核形状，
    4) 参数数量，
    5) 运算次数（乘加操作），
    6) 该层是否可训练。

注意：如果既没有提供 input_data 也没有提供 input_size，则不会执行前向传播，此时提供的模型信息仅限于层名称。

参数：
    model (nn.Module):
            要总结的 PyTorch 模型。模型应处于 train() 或 eval() 模式之一。如果各层模式不一致，运行 summary 可能会对 BatchNorm 或 Dropout 的统计产生副作用。若遇到此类问题，请提交 GitHub 问题。

    input_size (序列尺寸)：
            输入数据的形状，以 List/Tuple/torch.Size 形式提供。
            （数据类型需与模型输入匹配，默认为 FloatTensor）。
            元组中应包含批次大小。
            默认：None

    input_data (张量序列)：
            用于模型前向传播的参数（数据类型由输入推断）。
            如果 forward() 函数接受多个参数，请传入参数列表或关键字参数字典。
            （若 forward() 函数仅接受一个字典作为参数，需将其包裹在列表中）。
            默认：None

    batch_dim (int)：
            输入数据的批次维度。若 batch_dim 为 None，则假定 input_data / input_size 包含批次维度，并在所有计算中使用该维度。否则，会扩展所有张量以包含批次维度。
            指定 batch_dim 可以优化运行时性能，因为当指定了批次维度时，torchinfo 在前向传播时会使用批次大小为 1。
            默认：None

    cache_forward_pass (bool)：
            若为 True，则会以模型类名作为键缓存 forward() 函数的运行结果。如果前向传播操作开销较大，这将使修改模型摘要的格式更加方便，例如更改深度或启用的列类型，特别是在 Jupyter Notebook 中。
            警告：启用此功能后，若修改模型架构或输入数据/输入尺寸，缓存不会失效，也不会重新运行前向传播，从而可能导致摘要不准确。
            默认：False

    col_names (可迭代字符串)：
            指定在输出中显示哪些列。当前支持的列有：(
                "input_size",
                "output_size",
                "num_params",
                "params_percent",
                "kernel_size",
                "groups",
                "mult_adds",
                "trainable",
            )
            默认：("output_size", "num_params")
            如果未提供 input_data / input_size，则仅显示 "num_params"。

    col_width (int)：
            每列的宽度。
            默认：25

    depth (int)：
            显示嵌套层数的深度（例如 Sequential）。
            低于此深度的嵌套层将不会显示在摘要中。
            默认：3

    device (torch.Device)：
            使用指定的 torch 设备来处理模型和输入数据。
            如果未指定，则优先使用输入数据的数据类型（若有），否则使用模型参数的数据类型。若两者均未指定，则使用 torch.cuda.is_available() 的结果。
            默认：None

    dtypes (List[torch.dtype])：
            如果使用 input_size，torchinfo 假设输入数据为 FloatTensor 类型。若您的模型使用其他数据类型，请在此处指定相应的数据类型。
            对于多输入情况，需同时指定每个输入的尺寸及对应的数据类型。
            默认：None

    mode (str)：
            可取值为 "train"、"eval" 或 "same"，用于决定在调用 summary 之前是调用 model.train() 还是 model.eval()。无论哪种情况，最终都会恢复模型的原始模式。
            默认："same"。

    row_settings (可迭代字符串)：
            指定在每一行中显示哪些特性。当前支持的特性有：(
                "ascii_only",
                "depth",
                "var_names",
            )
            默认：("depth",)

    verbose (int)：
            0（静默）：无输出
            1（默认）：打印模型摘要
            2（详细）：完整显示权重和偏置层的信息
            默认：1
            如果使用 Jupyter Notebook 或 Google Colab，则默认为 0。

    **kwargs：
            `model.forward` 函数中使用的其他参数。不再支持 *args 的传递方式。

返回：
    ModelStatistics 对象
            更多信息请参阅 torchinfo/model_statistics.py 文件。
"""

示例

获取模型摘要的字符串表示

from torchinfo import summary

model_stats = summary(your_model, (1, 3, 28, 28), verbose=0)
summary_str = str(model_stats)
# summary_str 包含摘要的字符串表示！

探索不同的配置

class LSTMNet(nn.Module):
    def __init__(self, vocab_size=20, embed_dim=300, hidden_dim=512, num_layers=2):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers, batch_first=True)
        self.decoder = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        embed = self.embedding(x)
        out, hidden = self.encoder(embed)
        out = self.decoder(out)
        out = out.view(-1, out.size(2))
        return out, hidden

summary(
    LSTMNet(),
    (1, 100),
    dtypes=[torch.long],
    verbose=2,
    col_width=16,
    col_names=["kernel_size", "output_size", "num_params", "mult_adds"],
    row_settings=["var_names"],
)

========================================================================================================================
Layer (type (var_name))                  Kernel Shape         Output Shape         Param #              Mult-Adds
========================================================================================================================
LSTMNet (LSTMNet)                        --                   [100, 20]            --                   --
├─Embedding (embedding)                  --                   [1, 100, 300]        6,000                6,000
│    └─weight                            [300, 20]                                 └─6,000
├─LSTM (encoder)                         --                   [1, 100, 512]        3,768,320            376,832,000
│    └─weight_ih_l0                      [2048, 300]                               ├─614,400
│    └─weight_hh_l0                      [2048, 512]                               ├─1,048,576
│    └─bias_ih_l0                        [2048]                                    ├─2,048
│    └─bias_hh_l0                        [2048]                                    ├─2,048
│    └─weight_ih_l1                      [2048, 512]                               ├─1,048,576
│    └─weight_hh_l1                      [2048, 512]                               ├─1,048,576
│    └─bias_ih_l1                        [2048]                                    ├─2,048
│    └─bias_hh_l1                        [2048]                                    └─2,048
├─Linear (decoder)                       --                   [1, 100, 20]         10,260               10,260
│    └─weight                            [512, 20]                                 ├─10,240
│    └─bias                              [20]                                      └─20
========================================================================================================================
Total params: 3,784,580
Trainable params: 3,784,580
Non-trainable params: 0
Total mult-adds (M): 376.85
========================================================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.67
Params size (MB): 15.14
Estimated Total Size (MB): 15.80
========================================================================================================================

ResNet

import torchvision

model = torchvision.models.resnet152()
summary(model, (1, 3, 224, 224), depth=3)

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
ResNet                                   [1, 1000]                 --
├─Conv2d: 1-1                            [1, 64, 112, 112]         9,408
├─BatchNorm2d: 1-2                       [1, 64, 112, 112]         128
├─ReLU: 1-3                              [1, 64, 112, 112]         --
├─MaxPool2d: 1-4                         [1, 64, 56, 56]           --
├─Sequential: 1-5                        [1, 256, 56, 56]          --
│    └─Bottleneck: 2-1                   [1, 256, 56, 56]          --
│    │    └─Conv2d: 3-1                  [1, 64, 56, 56]           4,096
│    │    └─BatchNorm2d: 3-2             [1, 64, 56, 56]           128
│    │    └─ReLU: 3-3                    [1, 64, 56, 56]           --
│    │    └─Conv2d: 3-4                  [1, 64, 56, 56]           36,864
│    │    └─BatchNorm2d: 3-5             [1, 64, 56, 56]           128
│    │    └─ReLU: 3-6                    [1, 64, 56, 56]           --
│    │    └─Conv2d: 3-7                  [1, 256, 56, 56]          16,384
│    │    └─BatchNorm2d: 3-8             [1, 256, 56, 56]          512
│    │    └─Sequential: 3-9              [1, 256, 56, 56]          16,896
│    │    └─ReLU: 3-10                   [1, 256, 56, 56]           --
│    └─Bottleneck: 2-2                   [1, 256, 56, 56]          --

  ...
  ...
  ...

├─AdaptiveAvgPool2d: 1-9                 [1, 2048, 1, 1]           --
├─Linear: 1-10                           [1, 1000]                 2,049,000
==========================================================================================
Total params: 60,192,808
Trainable params: 60,192,808
Non-trainable params: 0
Total mult-adds (G): 11.51
==========================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 360.87
Params size (MB): 240.77
Estimated Total Size (MB): 602.25
==========================================================================================

多输入与不同数据类型

class MultipleInputNetDifferentDtypes(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1a = nn.Linear(300, 50)
        self.fc1b = nn.Linear(50, 10)

        self.fc2a = nn.Linear(300, 50)
        self.fc2b = nn.Linear(50, 10)

    def forward(self, x1, x2):
        x1 = F.relu(self.fc1a(x1))
        x1 = self.fc1b(x1)
        x2 = x2.type(torch.float)
        x2 = F.relu(self.fc2a(x2))
        x2 = self.fc2b(x2)
        x = torch.cat((x1, x2), 0)
        return F.log_softmax(x, dim=1)

summary(model, [(1, 300), (1, 300)], dtypes=[torch.float, torch.long])

或者，你也可以直接传入输入数据，torchinfo 会自动推断数据类型。

input_data = torch.randn(1, 300)
other_input_data = torch.randn(1, 300).long()
model = MultipleInputNetDifferentDtypes()

summary(model, input_data=[input_data, other_input_data, ...])

顺序容器与 ModuleList

class ContainerModule(nn.Module):

    def __init__(self):
        super().__init__()
        self._layers = nn.ModuleList()
        self._layers.append(nn.Linear(5, 5))
        self._layers.append(ContainerChildModule())
        self._layers.append(nn.Linear(5, 5))

    def forward(self, x):
        for layer in self._layers:
            x = layer(x)
        return x


class ContainerChildModule(nn.Module):

    def __init__(self):
        super().__init__()
        self._sequential = nn.Sequential(nn.Linear(5, 5), nn.Linear(5, 5))
        self._between = nn.Linear(5, 5)

    def forward(self, x):
        out = self._sequential(x)
        out = self._between(out)
        for l in self._sequential:
            out = l(out)

        out = self._sequential(x)
        for l in self._sequential:
            out = l(out)
        return out

summary(ContainerModule(), (1, 5))

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
ContainerModule                          [1, 5]                    --
├─ModuleList: 1-1                        --                        --
│    └─Linear: 2-1                       [1, 5]                    30
│    └─ContainerChildModule: 2-2         [1, 5]                    --
│    │    └─Sequential: 3-1              [1, 5]                    --
│    │    │    └─Linear: 4-1             [1, 5]                    30
│    │    │    └─Linear: 4-2             [1, 5]                    30
│    │    └─Linear: 3-2                  [1, 5]                    30
│    │    └─Sequential: 3-3              --                        (recursive)
│    │    │    └─Linear: 4-3             [1, 5]                    (recursive)
│    │    │    └─Linear: 4-4             [1, 5]                    (recursive)
│    │    └─Sequential: 3-4              [1, 5]                    (recursive)
│    │    │    └─Linear: 4-5             [1, 5]                    (recursive)
│    │    │    └─Linear: 4-6             [1, 5]                    (recursive)
│    │    │    └─Linear: 4-7             [1, 5]                    (recursive)
│    │    │    └─Linear: 4-8             [1, 5]                    (recursive)
│    └─Linear: 2-3                       [1, 5]                    30
==========================================================================================
Total params: 150
Trainable params: 150
Non-trainable params: 0
Total mult-adds (M): 0.00
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================

贡献

我们非常欢迎所有的问题和拉取请求！如果您想知道如何构建该项目：

• torchinfo 使用最新版本的 Python 进行开发。
  • 更改应向后兼容 Python 3.8，并遵循 Python 对旧版本的支持终止政策。
  • 运行 pip install -r requirements-dev.txt。我们使用所有开发依赖包的最新版本。
  • 运行 pre-commit install。
  • 若要使用自动格式化工具，运行 pre-commit run -a。
  • 若要运行单元测试，运行 pytest。
  • 若要更新预期输出文件，运行 pytest --overwrite。
  • 若要跳过输出文件测试，使用 pytest --no-output。

参考文献

• 感谢 @sksq96、@nmhkahn 和 @sangyx 为本项目提供了灵感。
• 关于模型大小估算，感谢 @jacobkimmel（详情请见此处）。

torchinfo 快速上手指南

torchinfo 是一个用于 PyTorch 模型可视化的工具，功能类似于 TensorFlow 的 model.summary()。它能清晰地展示模型的层级结构、输入/输出形状、参数量及计算量（Mult-Adds），是调试神经网络结构的得力助手。

环境准备

• Python 版本：3.8 及以上
• PyTorch 版本：1.4.0 及以上
• 系统要求：支持 Linux、macOS 和 Windows

安装步骤

推荐使用 pip 进行安装。国内用户可使用清华源或阿里源加速下载。

使用 pip 安装（推荐）：

pip install torchinfo

使用国内镜像源加速安装：

pip install torchinfo -i https://pypi.tuna.tsinghua.edu.cn/simple

使用 Conda 安装：

conda install -c conda-forge torchinfo

基本使用

1. 最简单的用法

只需导入 summary 函数，传入模型实例和输入数据的形状（需包含 Batch Size），即可打印模型摘要。

from torchinfo import summary
import torch.nn as nn

# 假设你有一个定义好的模型
class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 320)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = ConvNet()

# 生成模型摘要
# input_size 格式为: (Batch_Size, Channels, Height, Width)
summary(model, input_size=(16, 1, 28, 28))

输出示例：

================================================================================================================
Layer (type:depth-idx)          Input Shape          Output Shape         Param #            Mult-Adds
================================================================================================================
ConvNet                         [16, 1, 28, 28]      [16, 10]             --                 --
├─Conv2d: 1-1                   [16, 1, 28, 28]      [16, 10, 24, 24]     260                1,048,320
├─Conv2d: 1-2                   [16, 10, 12, 12]     [16, 20, 8, 8]       5,020              2,248,960
├─Linear: 1-3                   [16, 320]            [16, 50]             16,050             112,350
├─Linear: 1-4                   [16, 50]             [16, 10]             510                3,570
================================================================================================================
Total params: 21,840
Trainable params: 21,840
Non-trainable params: 0
Total mult-adds (M): 3.41
================================================================================================================

2. Jupyter Notebook / Colab 特别说明

如果在 Jupyter Notebook 或 Google Colab 中使用，确保 summary() 是单元格的最后一个表达式，或者将其包裹在 print() 中：

# 方式一：作为单元格最后一行
summary(model, input_size=(16, 1, 28, 28))

# 方式二：显式打印
print(summary(model, input_size=(16, 1, 28, 28)))

3. 获取摘要字符串

如果你需要将摘要信息保存为字符串而不是直接打印，可以设置 verbose=0：

model_stats = summary(model, input_size=(16, 1, 28, 28), verbose=0)
summary_str = str(model_stats)
# 现在 summary_str 包含了完整的摘要文本