做深度学习久了你会清楚地感觉到一个尴尬时刻模型在训练集上准确率明明还行一到要上线要么显存爆了要么单次推理耗时超标。这时候你最需要回答的问题就是“这个模型到底有多少参数量、跑一次要多少计算量”。今天这篇就专门聊清楚一件事——用PyTorch怎么快速统计模型的参数量Params和计算量FLOPs以及如何通过现成工具把模型结构看得明明白白。我会从最笨的print(model)一直讲到thop、ptflops、fvcore这些常用工具过程中会穿插手动计算公式和踩坑记录。不管你是在做论文复现、模型压缩、边缘部署还是只是好奇自己搭的CNN每天在忙活什么这篇文章都应该能帮上忙。1. 为什么要关心参数量和计算量1.1 参数量和FLOPs到底是什么参数量英文叫Params指的是模型里所有可学习的参数个数。这些参数包括卷积核的权重、全连接层的矩阵权重、批归一化层的缩放因子和偏置等。可以简单理解成模型“记忆容量”参数量越大模型能装下的“知识”就越多但模型文件也越大内存和显存占用也越高。计算量英文叫FLOPs是Floating Point Operations的缩写指模型前向推理一次要做的浮点运算次数。FLOPs更接近模型“计算速度”或“计算复杂度”的度量如果参数量是模型的“体重”那FLOPs就是模型全力跑一百米要“使出的劲儿”。注意这里有个很容易混淆的概念叫MACsMultiply-Accumulate Operations指乘加运算次数。一次乘加包含乘法和加法两步所以严格来说1个MACs约等于2个FLOPs。很多工具报告的数据其实单位是MACs但都习惯性写成FLOPs这就埋下了不少争论的伏笔后面我会专门讲。1.2 为什么部署场景下这两个数字决定生死几乎所有需要把模型真正用起来的场景都绕不开这两个指标。第一个场景是边缘设备部署。手机、摄像头、嵌入式板卡的内存和算力都很有限一个模型如果参数量有几千万光权重文件就上百兆别说运行加载都费劲。这时参数量就决定了模型能不能塞进目标设备。第二个场景是实时性判断。无人车、直播美颜、实时翻译这类低延迟任务通常要求在几十毫秒内完成一次推理。模型FLOPs过大即使参数量看着不大也可能因为计算密度太高而跑不满帧率。第三个场景是论文和项目汇报。发论文要做模型复杂度对比不同模型如果不统一计算口径对比结果会非常奇怪。我在实际处理过的一个“模型性能对比表”里见过别人用不同库算出的参数差异根源就是有人统计了Buffers比如BN的running_mean有人只统计了parameters。这就逼着我们必须搞懂这些工具背后的统计规则。2. 查看模型结构从print到summary2.1 最直接的 print(model) 能看出什么刚接触PyTorch时最顺手的工具就是直接print(model)。比如定义了一个简单的卷积模型打印出来会看到这样的结构import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) self.bn1 nn.BatchNorm2d(16) self.conv2 nn.Conv2d(16, 32, 3, padding1) self.bn2 nn.BatchNorm2d(32) self.fc nn.Linear(32 * 8 * 8, 10) def forward(self, x): x nn.functional.relu(self.bn1(self.conv1(x))) x nn.functional.max_pool2d(x, 2) x nn.functional.relu(self.bn2(self.conv2(x))) x nn.functional.max_pool2d(x, 2) x x.reshape(x.size(0), -1) return self.fc(x) model SimpleNet() print(model)输出会按模块嵌套关系打印出每一层但不会告诉你每一层输出特征图的尺寸变化也不会告诉你每层参数数量更不会统计总体参数量。这个方式适合快速确认模块嵌套关系比如自己有没有把卷积和BN顺序搞反但要写报告就显然不够用了。2.2 torchinfo.summary结构、输出形状、参数量一站搞定真正用得多的看结构工具是torchinfo。它提供了类似Keras里model.summary()的体验安装很简单pip install torchinfo然后运行from torchinfo import summary summary(model, input_size(1, 3, 32, 32), col_names[input_size, output_size, num_params, mult_adds], depth4)这样能看到每一层接收的输入尺寸、输出的尺寸、该层的参数量以及累计的乘加运算量mult_adds。我最常用它来排查维数写没写对因为在Summary输出里如果某个layers的输出shape和预想的不一样一眼就能发现是池化kernel算错还是reshape有问题。这个工具对多数CNN、MLP都有效。但如果你用的是带自定义操作的模型比如自己写的复杂注意力模块偶尔会遇到“Unable to infer”这种报错这时需要在summary里指定device或给模型加一个threshold参数实话说排查起来有点麻烦所以模型结构特别复杂时我通常只用它打印“是否有预期输出形状”而不是完全依赖。2.3 用named_modules自定义遍历结构还有一个灵活度更高的方式自己递归遍历模型的模块树。PyTorch里每个nn.Module都支持named_modules()和named_parameters()写一个小工具就能输出结构树def print_struct(model, show_paramsFalse): total 0 for name, module in model.named_modules(): param_count sum(p.numel() for p in module.parameters()) total param_count if module is not model else 0 if param_count 0 or isinstance(module, nn.Conv2d): print(f{name}: {module.__class__.__name__}, params{param_count}) print(fTotal params: {total})这种方式的好处是能完全控制输出内容。我在对比不同模型变体时会额外标记出哪些模块被冻结parameters().requires_gradFalse方便快速定位可训练参数量。要使用print_struct时注意避免重复计数named_modules()会遍历所有子模块所以通常按叶子节点统计或者只统计model.parameters()不然容易把父模块的参数量重复累加。3. 统计参数量三种最常用方法3.1 利用parameters()手动求和最不加修饰、也最可信的参数量统计方法就是遍历模型参数total_params sum(p.numel() for p in model.parameters()) print(fTotal params: {total_params}) trainable_params sum(p.numel() for p in model.parameters() if p.requires_grad) print(fTrainable params: {trainable_params})这个统计的是所有Parameter类型的张量。对于卷积层numel()是输出通道数 × 输入通道数 × 卷积核高 × 卷积核宽再加上偏置参数的个数。比如nn.Conv2d(3, 16, 3, padding1)如果带偏置参数量是3 * 16 * 3 * 3 16 448。遍历出来的结果既直观又不容易出错唯一的缺点是如果模型较大手动统计一遍会慢一点但对常规模型完全没问题。这里必须说一个容易踩的坑model.parameters()不包含BN层的running_mean和running_var因为它们不是Parameter而是Buffer。也就是说纯参数统计并不能反映模型在显存中实际占用的全部存储空间。如果你的模型里有BatchNorm用“参数量”当“模型大小”是不够的还得额外把buffer加起来。我曾经在恢复某个BN特别多的小型网络时就因为这个忽略了buffer导致到板子上加载时内存比预期多出一块。如果需要统计Buffer占用可以用model.buffers()total_buffer sum(b.numel() for b in model.buffers())然后模型实际权重状态一共可能是params buffers。写报告时建议说清楚“参数量”是多少、“Buffers”有多少避免其他人复现时对不上。3.2 torchinfo summary里的参数量怎么读才准确前面提到torchinfo.summary会输出Total params、Trainable params 和 Non-trainable params。这些数值其实内部也是通过遍历model.parameters()得到的所以和前一种手动统计的结果理论上完全一致。但有一个细节要注意如果你给summary传入的input_size和实际训练尺寸不一样Summary表里某些层的参数计算可能受“特征图尺寸”影响尤其是包含AdaptivePool或自定义层的模型。大多数情况下参数个数不会变但Mult-Adds会变因此需要确保input_size符合真实数据形状。用summary还有一个好处它可以直观展示每一层的参数分布。比如你怀疑某个全连接层参数太多看表格立刻就能定位。通常fc层因为要把前面的特征图展平成向量会出现参数量爆炸这也是论文里把全连接替换成全局平均池化的常见原因。3.3 thop.profile一边计算量一边算参数量thop是老牌工具安装pip install thop用法比较固定import torch from thop import profile input torch.randn(1, 3, 32, 32) model.eval() macs, params profile(model, inputs(input,)) print(fMACs: {macs / 1e6:.2f} M, Params: {params / 1e6:.2f} M)这里名字里第一项返回的是macs不要习惯性叫它FLOPs。如果你想得到传统意义上的FLOPs大约要乘2。很多网上的博客不求甚解地把这个数字直接标成FLOPs结果和其他工具不一致这是最常见的“口径混乱”来源之一。还要注意profile默认会统计模型的全部参数包括那些requires_gradFalse的层如果需要只统计可训练参数可以加trainable_onlyTrue参数。thop对常见卷积、池化、全连接都有预设的钩子效率比较高。但遇到复杂的自定义算子比如某些注意力实现里的矩阵乘法它可能统计得不准确甚至直接提示没有实现对应的hook。这时候我需要回到“手动算”的老思路或者用下面要讲的ptflops。3.4 用ptflops统计附带逐层报告如果觉得thop的输出太单薄可以试ptflopspip install ptflops用法from ptflops import get_model_complexity_info macs, params get_model_complexity_info(model, (3, 32, 32), as_stringsTrue, print_per_layer_statTrue) print(MACs:, macs) print(Params:, params)ptflops的好处是能输出逐层的MACs和params并且支持print_per_layer_statTrue时打印一份类似summary的统计表。它的统计逻辑是“计算图中出现的乘加运算”对标准卷积层统计得比较准确。如果你正在写模型对比报告建议固定用ptflops作为标准工具至少在同一个表格里保持一致。4. 计算FLOPs原理与工具4.1 FLOPs的计算口径为什么让人头大先别急着跑工具得先弄明白FLOPs到底怎么算。拿一个卷积层举例。假设输入特征图尺寸是H_in × W_in卷积核是K × K输入通道数是C_in输出通道数是C_out输出特征图尺寸是H_out × W_out。在不考虑偏置的情况下每输出一个特征图上的点需要做C_in × K × K次乘法以及对应的C_in × K × K次加法将各通道乘积累加所以单点运算量约是2 × C_in × K × K。整个输出特征图上有H_out × W_out × C_out个点于是卷积层的FLOPs可以粗略写成FLOPs ≈ 2 × C_in × K × K × C_out × H_out × W_out如果加了偏置每个输出点额外多一次加法数值上只增加一点点很多工具会忽略这一步。而大多数库返回的“MACs”只统计了乘法和加法中更核心的乘加对因此数值上约等于上面公式的一半也就是MACs ≈ C_in × K × K × C_out × H_out × W_out这就是为什么代码里的macs和理论FLOPs常常差一倍。BN层在推理阶段通常会被融合进去所以在很多统计工具里不单独计算BN的FLOPs但在训练阶段BN每层也需要浮点运算。激活函数如ReLU通常不计入FLOPs因为它的运算量相比卷积和矩阵乘来说很小。明白了这些口径差异后面看工具输出就不会一脸迷茫了。4.2 thop.profile快速上手与踩坑thop.profile的使用我在前面已经写了核心代码。再补充几个实际操作中的注意点。第一传入的inputs必须是元组或列表的形式哪怕只有一个输入也要写(input,)否则容易报错。第二如果模型里包含Dropout请务必提前切到model.eval()虽然profile在统计FLOPs时不一定关心随机掩码但遵循训练和推理一致的做法更稳妥。第三对于带if分支的模型比如不同输入shape走不同分支profile只会统计你传入的这条输入路径的计算量这不算bug只是“单样本路径”快照。我自己踩过的一个坑是profile在处理自定义torch.nn.functional.conv2d时因为早期版本不支持某些参数组合给出的FLOPs会明显少。后来我改用ptflops的get_model_complexity_info并用official参数指定官方实现问题就少了。4.3 ptflops的常用参数详解get_model_complexity_info的几个关键参数值得记一下input_res一般传(C, H, W)。注意这里不是batch维度默认batch是1。as_strings为True时直接返回“xxx M”这种字符串方便人类阅读为False时返回浮点数。print_per_layer_stat为True时输出逐层报告调试时建议开。input_constructor如果模型需要多个输入或者需要传入掩码、文本等就得自己写一个input_constructor函数。例如有两个输入的模型def input_constructor(input_res): image torch.randn(1, *input_res) mask torch.ones((1, 1, input_res[1], input_res[2])) return {x: image, mask: mask}这一点让ptflops比thop更灵活。4.4 用fvcore做更细粒度的统计Meta开源的fvcore也提供了FLOPs统计工具而且它是facebookresearch里很多研究的基石。安装pip install fvcore统计代码from fvcore.nn import FlopCountAnalysis, parameter_count_table model.eval() input torch.randn(1, 3, 32, 32) flop_analyzer FlopCountAnalysis(model, (input,)) print(fFLOPs (fvcore): {flop_analyzer.total()}) print(parameter_count_table(model))fvcore的统计结果是“浮点运算次数”不过它同样会把卷积层的乘加当作两次浮点运算吗这一点需要看底层实现。我的经验是不同库对“是否为乘加分别计数”的定义并不相同所以在写报告时最好标注“本表格使用xx库统计单位是MACs/FLOPs”而不是只贴数字。fvcore还支持给特定算子注册自定义统计函数比如你有一个手写的稀疏卷积算子可以手动添加hook覆盖默认统计。这已经是比较进阶的用法了但当你面对的是VAE、扩散模型这类结构时会非常有效。4.5 计算FLOPs前记得固定输入形状和模式无论是哪个库FLOPs都和输入尺寸强相关。同一模型处理224×224和处理32×32理论FLOPs会差几十倍甚至上百倍因为卷积输出特征图大小直接受输入影响。所以比较两个模型时务必使用相同的输入分辨率并且明确写出来。还要记得模型切换成评估模式model.eval()。虽然Dropout和BN在train模式下也有计算量但你网络发布后通常用的是eval模式统计这个更符合真实部署。另外with torch.no_grad()不会影响FLOPs统计因为工具统计的是计算图结构并不真正做梯度计算。5. 完整实例一个自定义CNN模型的统计实战5.1 搭建一个样例网络为了把上面的方法串起来我搭一个小型CNN输入输出都简单import torch import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(3, 16, 3, padding1) # 输入3通道输出16通道 self.bn1 nn.BatchNorm2d(16) self.pool1 nn.MaxPool2d(2) self.conv2 nn.Conv2d(16, 32, 3, padding1) # 输入16输出32 self.bn2 nn.BatchNorm2d(32) self.pool2 nn.MaxPool2d(2) self.fc nn.Linear(32 * 8 * 8, 10) def forward(self, x): x torch.relu(self.bn1(self.conv1(x))) x self.pool1(x) x torch.relu(self.bn2(self.conv2(x))) x self.pool2(x) x x.reshape(x.size(0), -1) return self.fc(x)这个模型的输入是1×3×32×32的图片经过两次池化后特征图变成32×8×8然后展平接全连接层。5.2 用torchinfo和print_struct看结构先运行model SimpleNet() summary(model, input_size(1, 3, 32, 32), col_names[output_size, num_params])输出表格里会显示conv1.output_size[1,16,32,32]params是448conv2.output_size[1,32,16,16]params是4672两个BN层各自参数32最后的fc层参数量是32*8*8 * 10 10 20480 10 20490。汇总一下总参数就是这些数字之和约26000左右。跑一下手动统计total_params sum(p.numel() for p in model.parameters()) print(total_params)结果没问题。如果想知道“实际模型状态”的总量还得加上BN里running_mean和running_var的buffer每个BN层两个buffer每个buffer长度16或32共额外2×(105)30个数字不对还是以实际工具输出为准。这类细节虽然小但在统计不到位时会带偏报告。5.3 用thop和ptflops对比结果实际跑一遍input_tensor torch.randn(1, 3, 32, 32) macs, params profile(model, inputs(input_tensor,)) print(fthop - MACs: {macs / 1e6:.2f} M, Params: {params / 1e6:.2f} M)对SimpleNet来说主要计算集中在两个卷积层和全连接层。第一层卷积输入输出都是32×32MACs大约是3×3×3×16×32×32 442368。第二层输入16×16输出16×16MACs是16×3×3×32×16×16 1179648。全连接层MACs是(32*8*8) × 10 20480。三个加一起大约164.6万MACs。再乘2就是约329万FLOPs。ptflops给出的MACs数字也应该落在同一量级因为单位换算后都是“1.6M MACs”左右。这里能清楚看到“MACs”和“FLOPs”是两个数值如果写作时把thop的macs直接写成flops你报告里可能少算一倍。我通常会在代码里一次性打印两个指标并且标明单位。5.4 手动验算一遍更安心验算卷积层参数和MACs我亲自算一遍给新入门的朋友看。以conv1为例输入通道C_in3输出通道C_out16卷积核K3输入大小32×32输出大小32×32因为padding1。参数量3 * 16 * 3 * 3 16 448MACs3 * 16 * 3 * 3 * 32 * 32 442368FLOPs含乘和加2 * 3 * 16 * 3 * 3 * 32 * 32 884736手动公式算出的数字和工具基本一致。之所以建议动手验算是因为这能帮你建立对工具输出的直觉感将来遇到输出特别异常时能快速判断是模型结构问题还是工具统计问题。6. 常见问题与排查技巧6.1 为什么不同人统计的参数量不一样这个问题我在帮别人对代码时遇到过很多回。根源大致有三个一是没有区分parameters和buffers二是使用了torch.save保存整个模型加载后统计属性不同三是用了DataParallel包装模型导致model.module前多了一层部分统计脚本没写好就会漏一层。正确做法是指定统计某个固定模型的叶子参数最好写成一个小脚本任何人拿同一份权重文件跑都得到相同数字。常见的一个误区是直接sum(x.numel() for x in model.parameters())这在大多数情况下没错。但遇到共享权重某些层参数指向相同对象时可能重复计数。所以如果你自己实现了嵌入层共享要对这种结构格外留意。6.2 FLOPs和MACs单位混用的解决办法单位不统一是最容易吵架的地方。“M”代表百万1e6“G”代表十亿1e9。有些库输出是Gitflops有些是GMacs。比如一个模型的MACs如果是1.2G MacsFLOPs大概就是2.4G Flops。如果换算错你写“模型需要2.4G Flops”其实可能是1.2G Flops差了整整一倍。我的习惯是在项目根目录放一个model_cost.py用统一函数统计输出格式固定为“MACs百万”和“FLOPs百万”并带上工具名。这样所有实验记录都能放在同一张表里对比不会因为今天用thop、明天用ptflops而变得混乱。6.3 Transformer以及动态分支模型怎么算Transformer不是传统CNN它更多是矩阵乘。self-attention的计算量主要由三部分组成生成Q、K、V的三个线性层计算注意力权重矩阵的矩阵乘法以及加权求和。以batch_size1为例序列长度是L隐藏维度是D那么QKV三个线性层每个大约是L×D×D次MACs三个就是3LDD。注意力得分矩阵计算是L×L×D次MACs加权求和也是L×L×D次MACs。所以整体复杂度有一个关于L的二次项这也是长序列慢的原因。thop和ptflops通常能识别标准的nn.Linear和torch.matmul但如果你在代码里自定义了复杂分支比如用一个if结构控制不同输入长度走不同算子工具只能统计当前输入路径的FLOPs。我在一个做自回归模型的项目里就发现因为忽略了循环部分统计出的FLOPs远低于真实推理。这种情况建议直接把公式写进代码手动统计关键算子的计算量再汇总到日志。6.4 把统计结果自动记录到日志或模型卡不要每次都从终端复制数字。更稳妥的方式是封装一个函数把统计结果自动写成JSON方便和模型文件放一起。参考写法def save_model_cost(model, input_res(3, 32, 32), file_namemodel_cost.json): from ptflops import get_model_complexity_info from torchsummary import summary import json macs_str, params_str get_model_complexity_info(model, input_res, as_stringsTrue, print_per_layer_statFalse) result { tool: ptflops, input_res: list(input_res), macs: macs_str, params: params_str } with open(file_name, w) as f: json.dump(result, f, indent2)以后每次训练完保存完权重就顺便调用一次这样在复盘模型演进时能一眼看到参数和计算量有没有突然暴增比翻聊天记录可靠多了。7. 一点个人实践经验最后聊点只有动手跑过才明白的经验。我在实际项目里发现工具统计出来的数字再准确也不如实机测试来得关键。尤其是算子融合、内存带宽限制、量化推理等优化之后FLOPs数字好看但实际速度依然不理想这是很常见的现象。所以我的工作流是先拿ptflops固定一个口径统计“静态复杂度”再跑到目标设备上实测一次延迟和吞吐两者一起看才不会被单一数字误导。如果你刚刚上手建议你现在就挑一个自己最近写的模型分别用torchinfo、thop、ptflops跑一遍再把我的示例里第二个卷积层手动算一次MACs对一下结果。脚踩这一圈之后你以后看任何模型复杂度报告都会心里有底。我个人还有一个特别喜欢的小技巧给model.eval()之后项目里所有模型统一调用一遍统计脚本把耗时和FLOPs一起打印在实验表格里既方便内部讨论也方便向别人解释优化效果。这个方法不需要什么高深技术但能省掉非常多“你这个数字怎么来的”的追问。
