PyTorchFLOPs计算与fvcore使用教程

本文深入解析了使用fvcore计算PyTorch模型FLOPs时常见不准甚至归零的根源——动态Python控制流、未注册自定义算子、输入张量缺失shape或device不匹配、jit.script模块被跳过，以及高层API get_model_flops的默认过滤陷阱；并给出切实可行的解决方案：统一用compute_flops配合model.eval()和torch.no_grad()、将if/else替换为torch.where等可trace操作、临时移除jit装饰器或内联展开、严格校验输入张量形态与设备，同时强调单位是flop（需除以1e9得GFLOPs），澄清4.1e9即4.1 GFLOPs而非4.1B的典型误解，助你精准评估模型计算复杂度，避开调试雷区。

fvcore.compute\_flops 为什么算不准自定义模型？

直接调用 fvcore.compute_flops 时返回 0 或远低于预期，大概率是模型里用了动态控制流（比如 if x.sum() > 0:）、未注册的自定义算子、或输入张量未带 shape 信息。fvcore 依赖 TorchScript 的 tracing，它只记录前向中实际执行的路径，不支持 symbolic shape 推导。

实操建议：

• 确保模型前向完全由标准 PyTorch 操作构成，避免 Python 控制流；若必须用，改用 torch.where 或 torch.nn.functional 等可 trace 的替代
• 传入的 input 必须是具体张量（不能是 torch.randn(1, 3, 224, 224, requires_grad=False) 这种没 detach 的），且 device 与模型一致（否则 tracing 失败）
• 对含子模块的模型，先用 model.eval() 和 torch.no_grad() 包裹，排除 dropout/batchnorm 统计更新干扰

怎么让 fvcore 正确识别带 torch.jit.script 的模块？

fvcore 默认不解析 @torch.jit.script 函数体，会把它当黑盒跳过，导致这部分 FLOPs 计为 0。这不是 bug，而是设计限制：jit script 函数无法被 Python AST 分析器读取。

解决方法只有两个：

• 临时移除 @torch.jit.script 装饰器，等 FLOPs 测完再加回去（适用于调试阶段）
• 把 jit 函数内联展开成普通 Python 实现，哪怕只在评估脚本里重写一遍（适合小函数）
• 若函数逻辑复杂且无法改写，可手动估算其核心循环次数 × 单次运算量，再加到 fvcore 总结果上

fvcore.get\_model\_flops 和 compute\_flops 有什么区别？

fvcore.get_model_flops 是高层封装，内部调用 compute_flops，但它默认启用 ignore_modules 过滤掉 torch.nn.AdaptiveAvgPool2d 等“零参数”但有计算的模块，容易低估；而 compute_flops 更底层、可控性更强。

推荐始终用 compute_flops，并显式指定参数：

from fvcore.nn import FlopCountAnalysis
flops = FlopCountAnalysis(model, inputs)
flops.unsupported_ops_warnings(False)  # 关闭警告干扰
flops.uncalled_modules_warnings(False)
print(flops.total())  # 单位是 flop（不是 GFLOPs）

注意：flops.total() 返回的是浮点运算次数，要转 GFLOPs 需除以 1e9；别误用 flops.by_module() 直接求和——它返回的是字典，sum(dict.values()) 才是总数。

为什么 ResNet50 算出来是 4.1G，但论文写 4.1B？

单位混淆是最常见的坑。FlopCountAnalysis.total() 返回的是 “flop”，即单次浮点运算（1 addition + 1 multiply = 2 flop，Conv2d 通常按 2×MACs 算）。ResNet50 输入 224×224 时理论值约 4.1×10⁹ flop → 4.1 GFLOPs，不是 GB 或 4.1B（后者是错把 10⁹ 当成 “Billion” 符号误读）。

验证方式很简单：

• 用 print(flops.total() / 1e9) 看是否 ≈ 4.1
• 对比 torchprofile 或 thop 输出，它们也返回 flop 数，单位一致
• 如果数值差 3 个数量级（比如得到 4.1e6），检查是否输错了 input shape（如把 (1,3,224,224) 写成 (1,3,24,24)）

真正容易被忽略的是：FLOPs 是理论峰值，不等于 GPU 实际耗时。内存带宽、kernel launch 开销、Tensor Core 利用率这些，fvcore 一概不建模。

终于介绍完啦！小伙伴们，这篇关于《PyTorchFLOPs计算与fvcore使用教程》的介绍应该让你收获多多了吧！欢迎大家收藏或分享给更多需要学习的朋友吧~golang学习网公众号也会发布文章相关知识，快来关注吧！