
做深度学习模型不管你是搞研究还是做工程落地都有一个绕不开的日常操作评估模型大小和计算开销。模型参数量Params决定文件体积和内存占用计算量FLOPs直接关乎推理速度和硬件选型。再加上时不时要查看模型每一层的结构细节这三件事基本构成了我调试模型的日常三大件。这篇博文就把这三件事一次说透怎么用PyTorch正确计算模型参数量、怎么统计FLOPs、怎么快速查看框架结构以及我在实际项目中踩过的坑和总结的避坑经验。适合正在做模型选型、神经网络调参、端侧部署评估或者写论文需要对比实验指标的读者。1. 为什么参数量和计算量是绕不开的两个指标很多刚入门的同学觉得训练一个模型能跑起来就完事了。实际上模型的参数量和计算量直接决定了这个模型能不能在你预期的硬件上跑起来跑多快功耗多高。1.1 参数量决定模型体积直接关联内存占用参数量统计的是模型里所有可学习的权重和偏置一共有多少个数值。假设一个模型参数量是100M按每个参数4字节的float32存储光权重文件就要占用约400MB空间。这个数字直接决定了模型文件能不能塞进嵌入式设备的Flash芯片GPU显存是否够用优化器和梯度还要额外占用通常按参数量的N倍估算端侧推理时内存带宽是否吃紧我在做端侧部署的时候有一个很实用的估算公式模型加载内存约为参数量乘4字节如果加上推理时的中间特征图显存需求通常是参数量乘以20到30倍。所以看到参数量心里基本就有数了。1.2 计算量决定推理速度影响硬件选型FLOPsFloating 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(pretrainedFalse) print(model)运行结果会以层级缩进的方式打印出完整的模块结构比如ResNet的conv1、bn1、layer1、layer2、fc等。这个方法虽然简单但有两方面不足一是每一层的输出张量形状看不到二是参数量要自己心里估算数字不会直接给出来。2.2 用summary查看每层输出形状和参数量比print更进一步的是借助torchsummary工具它能输出每一层的输出Shape、参数量并直接汇总总参数量。我经常用的是这类summary工具中的一个from torchsummary import summary model models.resnet18(pretrainedFalse).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, paramsdict(model.named_parameters())) dot.render(model_structure, formatpng)生成的图片会展示每一个操作节点以及数据流向。我的习惯是简单模型用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(pretrainedFalse) print(fTotal 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(fTrainable: {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每秒浮点运算次数是硬件性能指标。一个是数量一个是速率。MACsMultiply-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(pretrainedFalse) input_tensor torch.randn(1, 3, 224, 224) flops, params profile(model, inputs(input_tensor,)) flops, params clever_format([flops, params], %.3f) print(fFLOPs: {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(pretrainedFalse) macs, params get_model_complexity_info( model, (3, 224, 224), as_stringsTrue, print_per_layer_statTrue, ) print(fMACs: {macs}, Params: {params})注意ptflops默认输出的是MACs而不是FLOPs如果需要FLOPs可以手动乘以2。print_per_layer_statTrue会打印每一层的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), verboseTrue): model.eval() input_data torch.randn(*input_shape) flops, params profile(model, inputs(input_data,), verboseFalse) # 手动计算更准确的参数量 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×112FLOPs直接降到原来的四分之一左右但参数量不变。这带来的一个实用价值是在端侧部署时如果模型时延超标可以尝试降低输入分辨率这是性价比最高的优化手段。建议在统计FLOPs时顺带记录输入分辨率否则后续对比容易出现混淆。6.3 模型里有动态分支统计不准如果模型forward里有if条件根据输入情况走不同分支那么hook机制统计到的就只是实际走到的那条路径的计算量。比如一个输入尺寸可变的目标检测模型在不同尺寸输入下统计出的FLOPs就会有差异。解决方法是用和实际部署完全一致的输入尺寸和输入张量分布去做统计。如果模型存在输入分辨率自适应逻辑需要固定分辨率后统计或者分别统计不同尺度下的FLOPs标注清楚统计条件。我一般会在结果旁边加注释比如“FLOPs512×512”来避免误导。6.4 自定义算子无法识别怎么办如果用到了自定义的C扩展或自定义autograd.Functionthop和ptflops都无法识别统计结果会漏算这部分计算量。这时需要手动补充统计规则或者直接手算。thop提供一种自定义hook的机制但说实话接口不算友好。我更推荐的做法是把手写的自定义算子单独统计加到工具结果之上。比如一个自定义的相似度计算模块可以根据公式手算FLOPs然后和profile结果相加。6.5 CUDA环境不管用怎么办torchsummary在使用时如果模型在GPU上但显存不足或者机器上没有GPU就很麻烦。我的建议是统计这类指标时无需GPU直接用CPU模式即可。把模型放在CPU上输入张量也用CPU慢一点但不影响统计准确性。6.6 这个KI工具统计速度快慢有些同学反馈thop统计大模型比如7B参数的大模型语言模型时非常慢或者直接内存溢出。实际原因是profile前向传播过程中会真实跑一遍模型模型太大时中间激活值占用内存过多。解决办法是用更小的输入尺寸比如batch_size1不过多减小分辨率容易导致层维度报错用空输入或者只看一层的方式分段统计用更轻量的统计库直接基于参数配置计算比如根据模型配置流式统计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.2MFLOPs从1.2G降到310M精度只掉了0.3%。这个优化过程如果没有准确的FLOPs统计作为指导基本就是瞎调。这个例子想说明一件事统计指标不是拿来当报告看的而是用来指导模型设计和硬件匹配的。当你看到某一层占了40%的FLOPs就该考虑这一层是不是存在优化空间当你看到参数量集中在某个模块就该考虑是不是可以共享权重或者降维。做模型这么久我最大的感受是参数量和FLOPs这些数字不是你写完模型后补一个统计就完事的。最好在模型设计阶段就形成“结构设计、指标统计、瓶颈分析、结构改进”的循环习惯。每改一次结构顺手跑一次统计看着数字下降比训练完看精度曲线要更有手感因为你是在从源头控制模型的复杂度和效率。