torch_flops: 准确捕获forward中所有算子的FLOPs计算库
该文章介绍了torch_flops库,它是一个用于准确计算神经网络模型FLOPs(浮点操作次数)的库。在对比神经网络模型轻量程度时,FLOPs是一个重要的指标。然而,现有的FLOPs计算库存在一些问题,无法准确计算forward()中直接写出的运算符。torch_flops库通过使用pytorch的Symbolic Tracing功能,可以捕获forward中的所有算子,包括继承自nn.Module的子类、pytorch的函数运算以及tensor本身的方法,从而帮助用户更准确地评估神经网络模型的计算复杂度。
原创
2023-10-27 09:57:41 ·
386 阅读 ·
0 评论