做深度学习模型,不管你是搞研究还是做工程落地,都有一个绕不开的日常操作:评估模型大小和计算开销。模型参数量(Params)决定文件体积和内存占用,计算量(FLOPs)直接关乎推理速度和硬件选型。再加上时不时要查看模型每一层的结构细节,这三件事基本构成了我调试模型的日常三大件。
这篇博文就把这三件事一次说透:怎么用PyTorch正确计算模型参数量、怎么统计FLOPs、怎么快速查看框架结构,以及我在实际项目中踩过的坑和总结的避坑经验。适合正在做模型选型、神经网络调参、端侧部署评估,或者写论文需要对比实验指标的读者。
1. 为什么参数量和计算量是绕不开的两个指标
很多刚入门的同学觉得,训练一个模型能跑起来就完事了。实际上,模型的参数量和计算量直接决定了这个模型能不能在你预期的硬件上跑起来,跑多快,功耗多高。
1.1 参数量决定模型体积,直接关联内存占用
参数量统计的是模型里所有可学习的权重和偏置一共有多少个数值。假设一个模型参数量是100M,按每个参数4字节的float32存储,光权重文件就要占用约400MB空间。这个数字直接决定了:
- 模型文件能不能塞进嵌入式设备的Flash芯片
- GPU显存是否够用(优化器和梯度还要额外占用,通常按参数量的N倍估算)
- 端侧推理时内存带宽是否吃紧
我在做端侧部署的时候,有一个很实用的估算公式:模型加载内存约为参数量乘4字节,如果加上推理时的中间特征图,显存需求通常是参数量乘以20到30倍。所以看到参数量,心里基本就有数了。
1.2 计算量决定推理速度,影响硬件选型
FLOPs(Floating Point Operations,浮点运算次数)衡量的是模型做一次前向推理需要执行多少次浮点运算。同样一个模型,在不同硬件上的实际每秒浮点运算能力不同,所以通常用FLOPs去估算理论计算时间下限:
理论推理时间 ≈ 总FLOPs ÷ 硬件有效算力
举个例子,某个轻量模型在1TFLOPs算力的端侧芯片上,理论延迟就是总计算量除以芯片算力。有了FLOPs这个量化指标,你就不需要反复试跑不同硬件,直接可以估算出哪个模型适合哪块板子。
1.3 论文对比和算法选型中的硬指标
做算法调研的时候,你一定会看到论文里的参数量和计算量对比表格。同一类任务的不同模型,比如轻量级网络的不同实现,参数量和FLOPs往往是拉开差距的核心依据。同样是减少30%参数量,有的模型只损失0.5%精度,有的直接掉点2%,这就是选型取舍的关键。
说白了,参数量管模型“瘦不瘦”,FLOPs管模型“跑多快”,而查看框架结构则是搞清楚模型“长什么样”。这三件事互相配合,才是完整的模型评估体系。
2. 查看模型结构:从print到可视化
在计算指标之前,一般先要把模型结构摸清楚。PyTorch里查看模型结构有几层递进的方法,我平时都是混着用的。
2.1 最基础的print(model)输出
PyTorch里实例化模型后,直接print就能看到模型的模块嵌套关系:
import torchvision.models as models model = models.resnet18(pretrained=False) print(model)运行结果会以层级缩进的方式打印出完整的模块结构,比如ResNet的conv1、bn1、layer1、layer2、fc等。这个方法虽然简单,但有两方面不足:一是每一层的输出张量形状看不到,二是参数量要自己心里估算,数字不会直接给出来。
2.2 用summary查看每层输出形状和参数量
比print更进一步的是借助torchsummary工具,它能输出每一层的输出Shape、参数量,并直接汇总总参数量。我经常用的是这类summary工具中的一个:
from torchsummary import summary model = models.resnet18(pretrained=False).cuda() summary(model, input_size=(3, 224, 224))输出效果是这样的(示意):
---------------------------------------------------------------- Layer (type) Output Shape Param # ================================================================ Conv2d-1 [-1, 64, 112, 112] 9,408 BatchNorm2d-2 [-1, 64, 112, 112] 128 ReLU-3 [-1, 64, 112, 112] 0 ... ================================================================ Total params: 11,689,512 Trainable params: 11,689,512 Non-trainable params: 0 ----------------------------------------------------------------这个输出之所以重要,是因为在判断每一层是否存在维度匹配问题时非常直观。比如某个层输出维度和你预期不一致,在这个表里一眼就能找出来。需要注意:input_size这里的通道数要写全,不是224,而是(3, 224, 224),写错了会直接报错。
2.3 可视化方式:torchviz和TensorBoard
如果模型结构足够复杂,文本形式的打印看起来就有些吃力了。此时可以把计算图导出成图片,用torchviz的make_dot方法可以生成结构图:
from torchviz import make_dot x = torch.randn(1, 3, 224, 224).requires_grad_(True) y = model(x) dot = make_dot(y, params=dict(model.named_parameters())) dot.render("model_structure", format="png")生成的图片会展示每一个操作节点以及数据流向。我的习惯是:简单模型用print和summary足够,复杂模型或者需要写技术文档的时候才用make_dot生成图片,因为大模型的计算图会非常密集,肉眼辨认难度反而更大。
3. 参数量统计的两层理解
参数量计算在PyTorch里听起来很简单,就是遍历所有参数求和。但真正常用的统计方式,以及为什么不能只做一个简单求和,这里面的细节值得展开聊聊。
3.1 最朴素的统计方式:遍历parameters()
PyTorch中每个nn.Module都持有parameters()迭代器,返回的是所有叶子张量参数(包括权重和偏置)。最直接的统计代码:
def count_parameters(model): return sum(p.numel() for p in model.parameters()) model = models.resnet18(pretrained=False) print(f"Total parameters: {count_parameters(model)}")这里numel()是PyTorch中张量元素总数的函数,等价于shape各维度相乘。用这个方式统计ResNet18会得到11,689,512个参数,和torchsummary里的Total params一致。
3.2 为什么不能简单一刀切统计所有numel
上面这个朴素实现,在大多数模型上是准确的,但有几个特殊情况需要注意:
- 共享权重:如果模型里多个层引用了同一个参数对象(比如权重共享的Siamese网络),parameters()只返回一次该权重,所以统计是准确的。
- Buffer缓冲:BatchNorm里除了权重weight和偏置bias,还有running_mean和running_var这两个buffer。它们不参与梯度更新,但属于模型状态的一部分。用parameters()统计会漏掉buffer,而用state_dict()统计则会把它们算进去。
- 非叶子张量:如果一个参数是在forward过程中临时计算的,不参与反向传播的梯度更新,它也不会出现在parameters()里。
所以严谨的做法是区分“模型参数”和“模型状态”。一般而言,压缩包体积和显存占用估算应该以parameters()为准,但做模型量化或者状态迁移时则要考虑buffer。
3.3 可训练参数与全部参数的区分
很多模型存在“冻结层”的情况,比如迁移学习时冻结backbone,只训练分类头。这时如果笼统统计全部参数,会得到偏大的结果。PyTorch提供了requires_grad标记,可以分别统计可训练参数和冻结参数:
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) frozen_params = sum(p.numel() for p in model.parameters() if not p.requires_grad) print(f"Trainable: {trainable_params}, Frozen: {frozen_params}")这对端侧部署特别重要。很多部署框架只支持可训练参数的转换,冻结参数如果混在一起,转换时可能直接报错。我习惯在模型设计初期就打印一份可训练参数统计表,心里有个底。
4. 计算量FLOPs:概念辨析与统计工具
FLOPs这一块最容易懵,因为不同资料里的说法五花八门:FLOPs、FLOPs/s、MACs、MAdds。不把这几个概念理清楚,统计出来的数字意义也会打折扣。
4.1 FLOPs、MACs和FLOPS的区分
- FLOPs(小写s结尾):Floating Point Operations,浮点运算次数,是一个累计总量。我们常说的模型计算量就是这个。
- FLOPS(大写S结尾):Floating Point Operations per Second,每秒浮点运算次数,是硬件性能指标。一个是数量,一个是速率。
- MACs:Multiply-Accumulate Operations,乘加运算次数。一次MAC包含一次乘法和一次加法,这是实际硬件中最常见的操作单元。在很多加速器里,MACs比FLOPs更容易映射到硬件指令。
关键换算关系:一次卷积里的乘法和加法是分别计数的,一个乘加复合操作要算2次浮点运算。因此,在没有特殊说明的情况下,FLOPs ≈ 2 × MACs。而MAdds这个概念本质就是MACs的另一种叫法,意思是乘加次数。
4.2 卷积层计算量的直观推导
要理解工具输出的数字是怎么来的,最好自己动手推一次。以单个2D卷积为例,输出特征图每个位置的计算过程是:对输入通道数乘以卷积核高乘以卷积核宽的三维窗口做乘加操作。
计算公式可以写成:
FLOPs = 2 × C_out × H_out × W_out × (C_in × K_h × K_w + 1)
这里的分解逻辑是:
- C_out × H_out × W_out是输出元素个数
- C_in × K_h × K_w是每个输出元素做乘加的次数
- 加1是bias偏置的加法
- 乘以2是因为乘法算一次运算,加法算一次运算
比如输入是(64, 56, 56),用64个3×3卷积核输出也是(64, 56, 56),那么这一层FLOPs就是:
2 × 64 × 56 × 56 × (64 × 3 × 3 + 1) ≈ 72,413,184自己推导一遍的好处是,当统计工具报出可疑数字时,你可以快速手算做一个交叉验证。
4.3 常用统计工具对比
目前常用的FLOPs统计工具,我整理了下:
| 工具 | 特点 | 适用场景 |
|---|---|---|
| thop | 轻量,支持hook机制,老牌工具 | 快速统计,自定义扩展 |
| ptflops | 调用简单,直接传入input尺寸,输出整洁 | 日常实验对比 |
| torchprofile | 基于profiler机制,统计更贴近实际算子 | 排查算子级计算量 |
| fvcore | 来自某知名视觉库,支持多种API | 研究风格对比 |
这些工具的基本原理都是通过hook前向传播,在每一层执行时记录输入输出张量的shape,再根据层类型套用对应的计算量公式。也就是说,工具统计的FLOPs是理论估算值,不是实际硬件上的计算时间。
注意:不同工具对相同模型的统计结果可能有细微差异,主要是因为某些层(如ReLU、池化)是否计入了计算量的口径不同。论文对比时最好统一用同一个工具。
5. 实操:一行命令统计Params和FLOPs
说了这么多概念,现在进入正题。实际项目中我常用两种工具搭配使用,一个是thop,一个是ptflops,各有侧重。
5.1 基于thop的快速统计
thop的API设计很简洁,直接调用profile函数即可:
import torch from thop import profile, clever_format model = models.resnet18(pretrained=False) input_tensor = torch.randn(1, 3, 224, 224) flops, params = profile(model, inputs=(input_tensor,)) flops, params = clever_format([flops, params], "%.3f") print(f"FLOPs: {flops}, Params: {params}")这里需要注意的是一点:profile函数执行时会触发模型一次前向传播,所以模型要处于eval模式,否则BatchNorm等层会在前向过程中更新统计量。我建议在统计前加上model.eval()。
thop底层是通过注册hook来实现每一层计算量的累积统计的,它支持卷积、线性、BN、池化等常见层类型。如果你用了比较新的算子,thop可能不识别,需要手动注册统计规则。
5.2 基于ptflops的更直观用法
ptflops的调用方式和thop区别在于:不需要构造输入张量,直接传入input尺寸即可:
from ptflops import get_model_complexity_info model = models.resnet18(pretrained=False) macs, params = get_model_complexity_info( model, (3, 224, 224), as_strings=True, print_per_layer_stat=True, ) print(f"MACs: {macs}, Params: {params}")注意ptflops默认输出的是MACs而不是FLOPs,如果需要FLOPs可以手动乘以2。print_per_layer_stat=True会打印每一层的MACs和参数量,非常适合定位模型中的计算热点。
我自己习惯先把print_per_layer_stat打开看一遍,找出哪一层最重。很多模型其实就毁在几个高计算量层上,比如C3、C4阶段的3x3卷积,把这些层的通道数减半,FLOPs能直接降一大截。
5.3 手动排除BatchNorm层参数和计算量
在有些部署场景中,BatchNorm层在推理阶段会融合进前面的卷积层,所以统计时可以选择排除BN层。thop默认会统计BN的计算量,我个人用得比较多的自定义做法如下:
def count_params_without_bn(model): total = 0 for name, param in model.named_parameters(): if 'bn' in name or 'running' in name: continue total += param.numel() return total如果用thop想排除指定层类型,可以在profile之后手动减掉,或者用ignore_list参数。不过更稳妥的做法还是先原样统计,再结合自己的部署方案决定要不要排除BN。
5.4 把统计函数封装成通用工具
项目里如果经常要反复统计不同模型,每次都临时写一遍代码太浪费了。我都是封装成一个通用函数放在工具模块里,核心代码如下:
import torch from thop import profile def model_summary(model, input_shape=(1, 3, 224, 224), verbose=True): model.eval() input_data = torch.randn(*input_shape) flops, params = profile(model, inputs=(input_data,), verbose=False) # 手动计算更准确的参数量 all_params = sum(p.numel() for p in model.parameters()) trainable_params = sum( p.numel() for p in model.parameters() if p.requires_grad ) result = { "FLOPs": flops, "Params": all_params, "Trainable Params": trainable_params, } if verbose: for k, v in result.items(): if v > 1e9: print(f"{k}: {v / 1e9:.3f} G") elif v > 1e6: print(f"{k}: {v / 1e6:.3f} M") else: print(f"{k}: {v:.3f} K") return result这个封装把三种指标一次全部输出,平时我只要import这个函数,在模型训练入口调一行就能拿到完整报告,非常省事。
6. 常见问题与排查技巧实录
工具用多了,各种跑不通的情况也就见多了。我把这些年遇到的高频问题整理成速查表,方便你对号入座。
6.1 统计结果和论文对不上
这是最常遇到的坑。同一份模型结构,我用thop统计出来是2.5GFLOPs,但论文里写的是1.8GFLOPs,差了一大截。查了一圈发现原因有几种:
- 论文里的FLOPs可能输入分辨率不同(比如224 vs 256)
- 论文里的FLOPs可能没算bias、没算BN层
- 统计工具不同,对某些层是否计入的规则不同
- 论文里写的可能是MACs,但写作时误写了FLOPs
我的排查方法很简单:拿到论文的模型配置文件和代码,先用官方仓库的统计代码跑一遍,再和不同工具的结果交叉验证。如果官方用的输入分辨率是别的尺寸,统一按论文设置复现一次。
6.2 输入尺寸对FLOPs的影响被忽略
FLOPs和输入分辨率强相关。对于卷积层,输出特征图尺寸直接由输入尺寸决定。同样一个分类模型,输入从224×224改成112×112,FLOPs直接降到原来的四分之一左右,但参数量不变。
这带来的一个实用价值是:在端侧部署时,如果模型时延超标,可以尝试降低输入分辨率,这是性价比最高的优化手段。建议在统计FLOPs时,顺带记录输入分辨率,否则后续对比容易出现混淆。
6.3 模型里有动态分支统计不准
如果模型forward里有if条件,根据输入情况走不同分支,那么hook机制统计到的就只是实际走到的那条路径的计算量。比如一个输入尺寸可变的目标检测模型,在不同尺寸输入下统计出的FLOPs就会有差异。
解决方法是:用和实际部署完全一致的输入尺寸和输入张量分布去做统计。如果模型存在输入分辨率自适应逻辑,需要固定分辨率后统计,或者分别统计不同尺度下的FLOPs,标注清楚统计条件。我一般会在结果旁边加注释,比如“FLOPs@512×512”来避免误导。
6.4 自定义算子无法识别怎么办
如果用到了自定义的C++扩展或自定义autograd.Function,thop和ptflops都无法识别,统计结果会漏算这部分计算量。这时需要手动补充统计规则或者直接手算。
thop提供一种自定义hook的机制,但说实话接口不算友好。我更推荐的做法是:把手写的自定义算子单独统计,加到工具结果之上。比如一个自定义的相似度计算模块,可以根据公式手算FLOPs,然后和profile结果相加。
6.5 CUDA环境不管用怎么办
torchsummary在使用时如果模型在GPU上但显存不足,或者机器上没有GPU,就很麻烦。我的建议是统计这类指标时无需GPU,直接用CPU模式即可。把模型放在CPU上,输入张量也用CPU,慢一点但不影响统计准确性。
6.6 这个KI工具统计速度快慢
有些同学反馈thop统计大模型(比如7B参数的大模型语言模型)时非常慢,或者直接内存溢出。实际原因是profile前向传播过程中会真实跑一遍模型,模型太大时中间激活值占用内存过多。解决办法是:
- 用更小的输入尺寸(比如batch_size=1,不过多减小分辨率容易导致层维度报错)
- 用空输入或者只看一层的方式分段统计
- 用更轻量的统计库直接基于参数配置计算,比如根据模型配置流式统计
7. 一个端到端的小实践:边统计边调优
最后分享一个我在实际项目中习惯用的闭环流程。拿到一个新模型时,我不会直接去训练,而是先做一轮静态评估:
# 1. 查看结构 python show_structure.py # 2. 统计参数量和FLOPs python compute_flops.py --input-size 224 # 3. 根据统计结果判断是否需要精简模型有一次做扫码识别模型,初始结构用的标准分类网络,一统计发现FLOPs高达1.2G,在目标设备上理论延迟要50毫秒以上,完全超预算。我就用print_per_layer_stat逐层排查,发现重计算量集中在最后几个高维度卷积层上,于是在这些层加了深度可分离卷积替换,参数量从4.1M降到2.2M,FLOPs从1.2G降到310M,精度只掉了0.3%。这个优化过程如果没有准确的FLOPs统计作为指导,基本就是瞎调。
这个例子想说明一件事:统计指标不是拿来当报告看的,而是用来指导模型设计和硬件匹配的。当你看到某一层占了40%的FLOPs,就该考虑这一层是不是存在优化空间;当你看到参数量集中在某个模块,就该考虑是不是可以共享权重或者降维。
做模型这么久,我最大的感受是:参数量和FLOPs这些数字,不是你写完模型后补一个统计就完事的。最好在模型设计阶段就形成“结构设计、指标统计、瓶颈分析、结构改进”的循环习惯。每改一次结构,顺手跑一次统计,看着数字下降,比训练完看精度曲线要更有手感,因为你是在从源头控制模型的复杂度和效率。