Python怎么查看PyTorch模型的实际FLOPs_fvcore包结合网络结构进行计算量评估
fvcore.compute_flops算不准自定义模型,主因是动态控制流、未注册算子或输入缺shape;需避免Python控制流、确保输入为具体张量且device一致、用model.eval()和torch.no_grad();jit.script模块需临时移除装饰器或内联展开;推荐用compute_flops而非get_model_flops,并注意单位是flop(4.1e9=4.1 GFLOPs)。
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,并显式指定参数:
Python 3.14.3
微软官方的 Python 扩展,是 VS Code 安装量最高的扩展(209M+)。集成 IntelliSense(通过 Pylance)、调试(通过 Python Debugger)、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。
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 一概不建模。