1、基本介绍

torchinfo是一个为PyTorch用户量身定做的开源工具，其核心功能之一是summary函数。这个函数旨在简化模型的开发与调试流程，让模型架构一目了然。通过torchinfo的summary函数，用户可以快速获取模型的详细结构和统计信息，如模型的层次结构、输入/输出维度、参数数量、多加操作(Mult-Adds)等关键信息。

2、安装

首先，你需要安装torchinfo库。可以通过pip进行安装：

pip install torchinfo

3、导入

安装完成后，需要在你的Python脚本中导入torchinfo模块：

from torchinfo import summary

4、函数原型定义

torchinfo的summary函数原型定义如下：

def summary(model: nn.Module, input_data: torch.Tensor | tuple[torch.Tensor, ...] | tuple[int, ...] | None = None, batch_dim: int = 0, col_widths: tuple[int, ...] | None = None, col_names: tuple[str, ...] | None = None, device: str | torch.device | None = None, dtypes: tuple[torch.dtype, ...] | None = None, verbose: int = 1, **kwargs)

参数说明

• model: 要分析的PyTorch模型，必须是torch.nn.Module的实例。
• input_data: 用于模型前向传播的输入数据。它可以是一个torch.Tensor对象，也可以是一个包含多个输入张量的元组。此外，还可以提供一个表示输入尺寸的元组，例如(batch_size, channels, height, width)。
• batch_dim: 指定输入张量中哪个维度是批量大小（batch size）。默认为0。
• col_widths: 指定输出列宽的元组。如果未指定，则自动计算列宽以适应输出。
• col_names: 指定输出列名的元组。如果未指定，则使用默认列名。
• device: 指定模型运行的设备（如’cpu’或’cuda’）。如果未指定，则自动选择。
• dtypes: 指定输入张量的数据类型。如果未指定，则自动推断。
• verbose: 控制输出信息的详细程度。默认为1，表示输出基本信息。设置为2或更高可以获得更详细的输出。
• kwargs: 其他关键字参数，可以传递给模型的前向传播函数。

5、使用方法

下面通过几个示例来展示如何使用torchinfo的summary函数。
5.1 使用预定义模型
首先，我们使用PyTorch预定义的模型（如torchvision.models.resnet50）来展示如何使用summary函数。

import torch
import torchvision.models as models
from torchinfo import summary
# 定义模型
model = models.resnet18(pretrained=False)

# 使用summary函数打印模型概况
summary(model, input_size=(1, 3, 224, 224))

在这个示例中，我们加载了一个未预训练的ResNet50模型，并使用summary函数打印了模型的概况。input_size参数指定了输入数据的大小，即(batch_size, channels, height, width)。

5.2 使用自定义模型
接下来，我们定义一个简单的自定义模型，并使用summary函数打印其概况。

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

# 定义一个简单的两层全连接神经网络模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(100, 50)
        self.fc2 = nn.Linear(50, 10)
        self.relu = nn.ReLU()

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

# 创建模型实例
model = SimpleModel()

# 使用summary函数打印模型概况
summary(model, input_size=(100,))

在这个示例中，我们定义了一个简单的两层全连接神经网络模型，并使用summary函数打印了模型的概况。input_size参数指定了输入数据的大小，即(batch_size, features)。由于我们的模型是一个全连接层，所以我们只指定了特征数量。

5.3 使用自定义输入数据

有时候，可能想要使用实际的输入数据来查看模型的概况。下面是一个示例，展示了如何使用自定义输入数据来调用summary函数。

import torch
import torchvision.models as models
from torchinfo import summary

# 定义模型
model = models.resnet50(pretrained=False)

# 创建自定义输入数据
input_data = torch.randn(1, 3, 224, 224)  # batch_size=1, channels=3, height=224, width=224

# 使用summary函数打印模型概况
summary(model, input_data=input_data)

在这个示例中，我们创建了一个形状为(1, 3, 224, 224)的随机张量作为输入数据，并使用summary函数打印了模型的概况。注意，这里我们使用input_data参数而不是input_size参数来指定输入数据。

5.4 调整输出格式
torchinfo允许通过col_widths和col_names参数来调整输出的格式。下面是一个示例，展示了如何自定义输出列宽和列名。

import torch
import torchvision.models as models
from torchinfo import summary

# 定义模型
model = models.resnet50(pretrained=False)

# 使用summary函数打印模型概况，并自定义输出列宽和列名
summary(model, input_size=(3, 224, 224), col_widths=(30, 30, 20, 20),
        col_names=('input_size', 'output_size', 'kernel_size', 'num_params'))

在这个示例中，我们自定义了输出列宽和列名。col_widths参数指定了每列的宽度（以字符为单位），而col_names参数指定了每列的列名。这样，就可以根据需要来调整输出的格式了。

6、小结

torchinfo的summary函数是一个强大的工具，可以方便地查看PyTorch模型的结构和参数数量。通过本文的介绍，应该已经掌握了如何使用summary函数来打印模型的概况。无论使用预定义模型还是自定义模型，无论是使用输入尺寸还是自定义输入数据，torchinfo都能提供详细而清晰的输出信息。希望这篇文章能对你有所帮助！