从Keras转PyTorch？先搞定模型可视化：torchsummary vs torchinfo深度评测
当你从TensorFlow/Keras转向PyTorch时，最不习惯的可能就是模型结构的查看方式。Keras中简洁明了的model.summary()在PyTorch原生环境中并不存在，这让很多开发者感到困扰。本文将深入对比PyTorch生态中两个最流行的模型可视化工具——torchsummary和torchinfo，帮助你找到最适合的迁移方案。
1. 为什么PyTorch需要专门的模型可视化工具
PyTorch默认的print(model)输出虽然包含了所有层的信息，但存在几个明显问题：
import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 16, 3) self.pool = nn.MaxPool2d(2, 2) self.fc = nn.Linear(16 * 14 * 14, 10) def forward(self, x): x = self.pool(nn.functional.relu(self.conv1(x))) return self.fc(x.view(-1, 16 * 14 * 14)) model = SimpleCNN() print(model)输出结果：
SimpleCNN( (conv1): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1)) (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) (fc): Linear(in_features=3136, out_features=10, bias=True) )这种原生输出方式的主要缺陷包括：
- 缺乏层次结构：所有层平铺展示，难以理解网络架构
- 缺少关键信息：没有参数数量、输出形状等关键指标
- 可读性差：复杂模型会输出大量难以解析的信息
2. torchsummary：经典但有限的选择
torchsummary是最早解决这一问题的库之一，它的API设计参考了Keras的summary()。
2.1 基本使用
安装命令：
pip install torchsummary使用示例：
from torchsummary import summary summary(model, input_size=(3, 32, 32))典型输出：
---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 16, 30, 30] 448 MaxPool2d-2 [-1, 16, 15, 15] 0 Linear-3 [-1, 10] 313,610 ================================================================ Total params: 314,058 Trainable params: 314,058 Non-trainable params: 0 ---------------------------------------------------------------- Input size (MB): 0.01 Forward/backward pass size (MB): 0.11 Params size (MB): 1.20 Estimated Total Size (MB): 1.32 ----------------------------------------------------------------2.2 优缺点分析
优势：
- 输出格式接近Keras风格，迁移学习成本低
- 包含参数统计和内存估算
- 轻量级，依赖少
局限性：
- 不支持多输入/输出模型
- 自定义层处理能力有限
- 显存估算不够精确
- 已停止维护（最后一次更新在2019年）
提示：对于简单模型和快速原型设计，torchsummary仍然是一个不错的选择，但在复杂项目中可能会遇到限制。
3. torchinfo：现代PyTorch开发者的首选
torchinfo是torchsummary的现代替代品，提供了更全面的功能。
3.1 安装与基础使用
安装方式：
pip install torchinfo基本示例：
from torchinfo import summary summary(model, input_size=(1, 3, 32, 32))输出示例：
========================================================================================== Layer (type:depth-idx) Output Shape Param # ========================================================================================== SimpleCNN [1, 10] -- ├─Conv2d: 1-1 [1, 16, 30, 30] 448 ├─MaxPool2d: 1-2 [1, 16, 15, 15] -- ├─Linear: 1-3 [1, 10] 313,610 ========================================================================================== Total params: 314,058 Trainable params: 314,058 Non-trainable params: 0 Total mult-adds (M): 9.41 ========================================================================================== Input size (MB): 0.01 Forward/backward pass size (MB): 0.11 Params size (MB): 1.20 Estimated Total Size (MB): 1.32 ==========================================================================================3.2 高级功能
多输入支持：
class MultiInputModel(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(3, 16, 3) self.fc = nn.Linear(16 * 30 * 30 + 10, 10) # 额外10维输入 def forward(self, img, vec): x = nn.functional.relu(self.conv(img)) combined = torch.cat([x.view(-1, 16*30*30), vec], dim=1) return self.fc(combined) model = MultiInputModel() summary(model, input_data=[(1, 3, 32, 32), (1, 10)])自定义层可视化：
class CustomLayer(nn.Module): def __init__(self): super().__init__() self.weight = nn.Parameter(torch.randn(10, 10)) def forward(self, x): return x @ self.weight model = nn.Sequential( nn.Linear(20, 10), CustomLayer(), nn.ReLU() ) summary(model, input_size=(1, 20))深度信息展示： 通过depth参数控制显示层级：
summary(model, input_size=(1, 20), depth=3)3.3 核心优势对比
| 特性 | torchsummary | torchinfo |
|---|---|---|
| 多输入/输出支持 | ❌ | ✅ |
| 自定义层处理 | 有限 | 优秀 |
| 显存估算精度 | 一般 | 高 |
| 维护状态 | 停止 | 活跃 |
| 输出格式可定制性 | 低 | 高 |
| 计算量统计(MAdds) | ❌ | ✅ |
| 嵌套模型展示 | 平铺 | 树形 |
4. 实战建议与迁移策略
4.1 从Keras迁移的最佳实践
安装选择：
# 推荐使用torchinfo pip install torchinfo # 如果必须使用torchsummary pip install torchsummary等效代码对比：
Keras:
model = keras.Sequential([ layers.Dense(64, activation='relu'), layers.Dense(10) ]) model.build(input_shape=(None, 32)) model.summary()PyTorch等效:
model = nn.Sequential( nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 10) ) summary(model, input_size=(1, 32)) # torchinfo版本复杂模型处理技巧：
对于包含分支的结构，torchinfo能更好地展示：
class BranchModel(nn.Module): def __init__(self): super().__init__() self.conv = nn.Conv2d(3, 16, 3) self.branch1 = nn.Sequential( nn.Linear(16*30*30, 64), nn.ReLU() ) self.branch2 = nn.Sequential( nn.Linear(16*30*30, 32), nn.ReLU() ) def forward(self, x): x = nn.functional.relu(self.conv(x)) x = x.view(-1, 16*30*30) return self.branch1(x), self.branch2(x) model = BranchModel() summary(model, input_size=(1, 3, 32, 32))
4.2 性能考量
当模型特别大时，可以限制输出的详细程度：
# 只显示顶层模块 summary(model, input_size=(1, 3, 256, 256), depth=1) # 仅统计参数，不计算内存 summary(model, input_size=(1, 3, 256, 256), verbose=0)对于生产环境，可以考虑添加设备信息：
summary(model, input_size=(1, 3, 224, 224), device="cuda")4.3 调试技巧
遇到问题时，可以尝试以下方法：
形状不匹配：
# 使用colab或jupyter的调试模式 %debug # 或者逐层检查 for name, layer in model.named_children(): print(f"Testing layer: {name}") test_output = layer(test_input) print(f"Output shape: {test_output.shape}") test_input = test_output自定义层统计： 如果torchinfo无法正确统计你的自定义层，可以实现
extra_repr方法：class CustomLayer(nn.Module): def __init__(self): super().__init__() self.weights = nn.Parameter(torch.randn(10, 10)) self.bias = nn.Parameter(torch.zeros(10)) def forward(self, x): return x @ self.weights + self.bias def extra_repr(self): return f"weights={tuple(self.weights.shape)}, bias={tuple(self.bias.shape)}"
在实际项目中，我发现torchinfo几乎能覆盖所有模型可视化需求，特别是在处理Transformer等复杂架构时，其层次化展示方式比torchsummary直观得多。对于从Keras转来的开发者，建议直接使用torchinfo，它的输出格式更接近Keras的体验，同时提供了PyTorch特有的详细信息。
