3、PyTorch模型结构解析：学会使用torchsummary、torchinfo等工具查看模型参数、FLOPs与层结构，为压缩做准备
在Edge AI训练侧压缩工具链中，无论是结构化剪枝、非结构化剪枝、量化感知训练还是知识蒸馏，第一步都必须是“看清模型的真实结构”。很多工程师在拿到一个预训练模型后，直接套用压缩算法，结果往往因为对模型内部张量流动、参数分布、计算瓶颈缺乏量化认知，导致压缩后精度骤降或推理加速不明显。本章将系统讲解如何使用torchsummary、torchinfo等工具，从参数规模、计算量（FLOPs）、层间连接关系三个维度对PyTorch模型进行“解剖”，为后续压缩策略的制定提供数据基础。
3.1 核心原理：为什么压缩前必须先做结构解析
模型压缩的本质是在“冗余”与“精度”之间寻找平衡点。而冗余的识别，必须建立在精确的层级别统计之上。torchsummary和torchinfo的核心原理并不复杂：它们通过向模型输入一个虚拟的、符合实际推理场景的shape张量，触发一次前向传播，在传播过程中利用PyTorch的hook机制（register_forward_hook）拦截每一层的输入输出张量，从而计算出该层的输出特征图尺寸、参数量以及FLOPs。
这里需要特别强调的是，FLOPs（Floating Point Operations）与MACs（Multiply-Accumulate Operations）是两个容易混淆的概念。FLOPs通常指浮点运算次数，一次乘加运算（如y = w*x + b）在严格定义下计为2次FLOPs（1次乘法+1次加法）；而MACs只计1次乘加。torchinfo默认输出的是MACs，而部分和硬件厂商使用FLOPs。在对比不同模型的算力需求时，必须统一口径，否则会出现2倍的误差。
另一个核心原理是“输入分辨率敏感性”。模型的参数量是固定的，但FLOPs会随着输入张量的H和W变化而线性增长。例如一个Conv2d(3, 64, kernel_size=3, stride=1)，当输入为224x224时，其FLOPs约为224*224*64*3*3*3 ≈ 86.9M MACs；当输入为112x112时，FLOPs降为原来的1/4。因此，在解析模型时，必须使用目标部署场景的真实输入尺寸，否则得到的FLOPs参考值毫无意义。
关键认知：结构解析不是“看一眼模型有几层”，而是建立“层-参数-计算量-张量形状”的四维映射表。这张表是后续剪枝时判断哪些层是冗余的、量化时估计激活值范围、蒸馏时选择教师模型师生特征对齐位置的直接依据。
3.2 判断方法与操作流程：从安装到输出解读
目前社区最常用的工具是torchinfo（torchsummary的增强版），它修复了旧版torchsummary对多输入、复杂分支支持不佳的问题，并提供了更友好的表格输出。以下是标准操作流程：
- 安装工具库：
pip install torchinfo，同时确保PyTorch版本≥1.8。 - 构建模型实例：加载预训练权重或随机初始化模型，但必须保证模型处于
eval()模式，避免Dropout和BatchNorm在统计时产生干扰。 - 调用summary函数：传入模型、输入张量的shape（不需要真实数据，只需维度信息），并设置
batch_dim参数（通常为0）。 - 解读输出表格：关注每一行的
Output Shape、Param #、Mult-Adds列。
from torchinfo import summary
import torchvision.models as models
model = models.resnet18(pretrained=False)
# 关键：batch_size=1，输入通道3，分辨率224x224
summary(model, input_size=(1, 3, 224, 224),
col_names=["input_size", "output_size", "num_params", "mult_adds"],
row_settings=["var_names"])
输出表格中，Mult-Adds列即为MACs值。对于ResNet18，在224x224输入下，总MACs约为1.82G，总参数量约为11.69M。但请注意，这个数值包含了所有卷积层和全连接层，而最后的全连接层（fc层）参数量为512*1000=512K，占比约4.4%。如果部署目标是MCU级别的设备，这个全连接层可能就是需要被替换或剪枝的对象。
| 工具名称 | 输出内容 | 优点 | 局限性 |
|---|---|---|---|
| torchsummary（旧版） | 层类型、输出shape、参数量 | 轻量、简单 | 不支持多输入、无法显示FLOPs、对自定义模块支持差 |
| torchinfo（推荐） | 层类型、输入/输出shape、参数量、MACs、核心数 | 支持复杂模型、可自定义列、显示更全面 | 需要手动指定输入尺寸，对动态shape支持有限 |
| ptflops | FLOPs、参数量 | 轻量、快速 | 输出信息较少，无法查看逐层结构 |
| thop | FLOPs、参数量 | 兼容性好 | 对某些算子（如自定义op）统计不准确 |
3.3 案例拆解：MobileNetV2的结构解析与压缩点识别
我们以一个实际的Edge AI场景为例：需要将MobileNetV2部署到一颗带有NPU的SoC上，目标帧率≥30FPS，内存占用≤200MB。在压缩之前，首先进行结构解析。
初始现象与目标：原始MobileNetV2在GPU上运行流畅，但交叉编译到NPU后帧率仅12FPS，内存占用350MB。我们需要找出计算瓶颈和存储瓶颈。
关键输入数据：使用torchinfo对MobileNetV2进行解析，输入尺寸设为(1, 3, 224, 224)。解析结果摘要如下：
- 总参数量：3.50M（其中第18层卷积（1x1扩展层）参数量最大，约1.4M）
- 总MACs：300.8M（其中第18层和第14层贡献了约40%的MACs）
- 最大激活张量：第14层输出shape为(1, 96, 28, 28)，占用约300KB内存
分析判断过程：通过逐层查看Mult-Adds列，发现计算量集中在最后几个倒残差块（bottleneck），这是因为这些层的通道数较多（扩展因子为6）。而参数量最大的层是最后的1x1卷积，它负责将高维特征映射到低维。对于NPU部署，1x1卷积虽然计算密度高，但内存访问频繁，容易成为带宽瓶颈。
实际操作步骤：
- 使用
summary导出逐层数据，保存为CSV文件。 - 按
Mult-Adds降序排列，找出前5个计算热点层。 - 按
num_params降序排列，找出前5个参数密集层。 - 对比两组列表，发现第18层（1x1卷积）同时出现在两个列表中，确定为优先剪枝对象。
结果对比：对第18层进行通道剪枝（剪掉30%通道），微调后模型MACs降至210M（降低30%），参数量降至2.8M，NPU实测帧率提升至22FPS，精度仅下降0.8%。进一步结合量化（INT8），最终帧率达到31FPS，内存占用降至180MB，满足部署要求。
经验复盘：结构解析的价值不在于“看”，而在于“比”。将逐层数据导出后，与硬件算力表（如NPU的MACs/cycle、DDR带宽）进行对比，才能精准定位瓶颈。很多工程师只关注总FLOPs，忽略了内存访问量（MAUC），导致剪枝后计算量下降但延迟不变，原因就是剪枝的层不是内存瓶颈层。
3.4 常见误区与总结复盘
在实际教学和项目评审中，我发现工程师在使用结构解析工具时经常陷入以下五个误区，这些误区会直接导致后续压缩策略的失效。
误区一：忽略BatchNorm层和偏置项的参数统计。torchinfo默认会统计BN层的参数（gamma和beta），这部分参数虽然计算量小，但在量化时需要特殊处理。如果压缩目标是极低比特量化（如2bit），BN层的参数分布会直接影响量化尺度。
误区二：使用错误的FLOPs口径。部分工具输出的是MACs，部分输出的是FLOPs（2倍关系）。在向硬件团队提需求时，必须明确说明是哪种口径，否则会导致算力评估偏差一倍。
误区三：对动态shape模型强行使用summary。如果模型包含nn.LSTM或根据输入长度变化的循环层，torchinfo无法准确统计。此时应改用torch.profiler进行实际运行时的FLOPs统计。
误区四：只看参数量，不看激活值内存。在Edge设备上，激活值内存往往比权重内存更关键，因为它与batch_size和输入分辨率直接相关。结构解析时，必须同时关注每一层输出张量的shape，估算峰值内存。
误区五：认为“层数越多，模型越大”。实际上，深度可分离卷积的参数量远小于标准卷积，但FLOPs可能更高（因为深度卷积的MACs效率低）。解析时不能凭直觉判断，必须以数据为准。
总结复盘：PyTorch模型结构解析是压缩工具链的“眼睛”。通过torchinfo等工具，我们能够建立模型的“体检报告”——哪些层是参数大户、哪些层是计算热点、哪些层是内存瓶颈。这份报告直接决定了后续选择剪枝、量化还是蒸馏策略。建议在实际项目中，将结构解析脚本固化到模型训练流程中，每次训练完成后自动生成一份JSON格式的解析报告，并纳入版本管理。这样当模型迭代时，压缩工程师可以快速对比不同版本之间的结构变化，避免重复劳动。
最后需要强调的是，结构解析不是一次性的工作。在剪枝过程中，每剪掉一层或若干通道后，都需要重新运行summary，确认MACs和参数量的实际变化是否符合预期。因为某些剪枝操作（如对shortcut连接层的剪枝）会导致后续层的输入通道数变化，进而影响FLOPs的计算。只有形成“解析→压缩→再解析”的闭环，才能确保压缩过程的可控性和可回溯性。