从TensorFlow转PyTorch？手把手教你用torchinfo实现Keras式model.summary()

从TensorFlow转PyTorch？用torchinfo实现Keras式模型摘要的完整指南

当你从TensorFlow/Keras转向PyTorch时，最怀念的功能之一可能就是那个简洁明了的model.summary()。在调试复杂网络时，能够一目了然地看到每层的输出形状、参数数量等信息简直是开发者的福音。而PyTorch原生的print(model)输出往往让人眼花缭乱，特别是面对深度网络时。这就是为什么torchinfo会成为PyTorch生态中如此重要的工具——它完美复现了Keras的模型摘要体验，甚至在某些方面做得更好。

1. 为什么PyTorch开发者需要torchinfo

PyTorch以其动态计算图和Pythonic的设计哲学赢得了大量开发者的青睐，但在模型可视化方面，它确实没有提供像Keras那样开箱即用的友好体验。原生的print(model)输出存在几个明显痛点：

• 信息组织混乱：嵌套的模块结构使得关键信息难以快速定位
• 缺少重要指标：没有直观的参数总量统计和可训练参数占比
• 输出形状不明确：无法直接看到各层的输出维度变化
• 内存占用未知：缺乏对模型内存占用的估算

torchinfo解决了所有这些痛点，它提供的摘要信息包括：

========================================================================================== Layer (type:depth-idx) Output Shape Param # ========================================================================================== ├─Sequential: 1-1 [64, 16, 16, 16] -- │ └─Conv2d: 2-1 [64, 16, 32, 32] 448 │ └─BatchNorm2d: 2-2 [64, 16, 32, 32] 32 │ └─MaxPool2d: 2-3 [64, 16, 16, 16] -- │ └─ReLU: 2-4 [64, 16, 16, 16] -- ... ========================================================================================== Total params: 34,168 Trainable params: 34,168 Non-trainable params: 0 Total mult-adds (M): 181.82 ========================================================================================== Input size (MB): 0.79 Forward/backward pass size (MB): 29.37 Params size (MB): 0.14 Estimated Total Size (MB): 30.29 ==========================================================================================

这种结构化输出让模型调试效率提升了数倍，特别是当你需要：

• 快速验证网络结构的正确性
• 分析各层的参数分布
• 估算模型的内存占用
• 比较不同架构的设计差异

2. torchinfo的安装与基础使用

2.1 安装方法

安装torchinfo非常简单，可以通过pip或conda完成：

# pip安装 pip install torchinfo # conda安装 conda install -c conda-forge torchinfo

注意：建议使用Python 3.7及以上版本，并与你的PyTorch版本保持兼容

2.2 基本用法

使用torchinfo.summary()函数生成模型摘要，其核心参数包括：

• model: 要分析的PyTorch模型实例
• input_size: 输入张量的形状（批处理大小需明确）
• depth: 显示的嵌套深度（默认为3）
• verbose: 控制输出详细程度

基础示例：

from torchvision.models import resnet18 from torchinfo import summary model = resnet18() summary(model, input_size=(1, 3, 224, 224))

对于更复杂的模型，你可能需要指定多个输入的形状：

# 多输入模型示例 summary( multi_input_model, input_data=[(1, 3, 256, 256), (1, 10)], # 两个输入的形状 dtypes=[torch.float32, torch.long] # 各自的数据类型 )

2.3 输出解读

torchinfo生成的摘要包含几个关键部分：

1. 层结构树：展示各层的类型、深度索引和输出形状
2. 参数统计：每层的可训练参数数量
3. 总量统计：
   • 总参数数量（区分可训练与不可训练）
   • 乘加操作总量（MAdds）
4. 内存估算：
   • 输入大小
   • 前向/反向传播中间变量大小
   • 参数存储大小
   • 总预估内存占用

这些信息对于模型优化和调试至关重要。例如，通过观察"Forward/backward pass size"可以识别内存瓶颈层，而"Total mult-adds"则反映了计算复杂度。

3. 高级功能与定制选项

3.1 自定义显示深度

对于特别深的网络（如ResNet152），你可能需要控制显示的层级深度：

# 只显示前5层细节 summary(model, input_size=(1, 3, 224, 224), depth=5) # 显示完整细节（可能非常长） summary(model, input_size=(1, 3, 224, 224), depth=10)

3.2 设备与数据类型支持

torchinfo可以正确处理不同设备和数据类型：

# GPU模型分析 model = model.to('cuda') summary(model, input_size=(1, 3, 224, 224), device='cuda') # 混合精度训练模型 with torch.cuda.amp.autocast(): summary(model, input_size=(1, 3, 224, 224))

3.3 批处理维度处理

torchinfo会自动处理批处理维度，但有时需要特别关注：

# 可变批处理大小分析 summary(model, input_size=(None, 3, 224, 224)) # 批处理维度可变 # 实际数据形状分析 batch_data = torch.randn(16, 3, 224, 224) summary(model, input_data=batch_data)

3.4 自定义列显示

你可以通过col_names参数定制显示的列：

summary( model, input_size=(1, 3, 224, 224), col_names=[ "input_size", "output_size", "num_params", "params_percent", "kernel_size", "mult_adds", ], )

可用列名包括：

• "input_size": 输入形状
• "output_size": 输出形状
• "num_params": 参数数量
• "kernel_size": 卷积核大小
• "mult_adds": 乘加操作数
• "trainable": 是否可训练

4. 与Keras model.summary()的深度对比

虽然torchinfo提供了类似Keras的摘要功能，但两者在实现和功能上存在一些差异：

Keras风格输出示例：

_________________________________________________________________ Layer (type) Output Shape Param # ================================================================= conv2d_1 (Conv2D) (None, 32, 32, 32) 896 _________________________________________________________________ batch_normalization_1 (Batch (None, 32, 32, 32) 128 _________________________________________________________________ max_pooling2d_1 (MaxPooling2 (None, 16, 16, 32) 0 ================================================================= Total params: 1,024 Trainable params: 960 Non-trainable params: 64 _________________________________________________________________

PyTorch torchinfo输出优势：

1. 内存分析：直接估算训练时所需内存，避免OOM错误
2. 计算量统计：MAdds指标帮助评估模型计算复杂度
3. 层级关系：树状结构更清晰反映模块嵌套关系
4. 灵活性：支持更多自定义选项和复杂模型结构

在实际项目中，我发现torchinfo的内存估算特别有用。例如，当处理大图像输入时，它能提前预警潜在的内存问题，这是原生的Keras摘要所不具备的。