Pytorch中计算自己模型的FLOPs(profile(),clever_format())

提到计算量必然涉及FLOPS, FLOPs和MACs，并不是同样的意思，切忌混淆。

FLOPS：Floating Point Operations Per Second，每秒浮点运算次数，是一个衡量硬件计算性能的指标。

FLOPs：Floating Point Operations，浮点运算次数，常用来衡量模型的计算复杂度，即计算量。（大多数论文会在表格中给出模型的FLOPs指标。不过，该指标与实际速度并不完全正相关，我自己测也是这样，所以也有一些论文会给FPS或者1/FPS即单图片推理时间。）

显然，FLOPS和FLOPs不是一回事，FLOPS的积分可以理解为等于FLOPs。

MACs：Multiply–Accumulate Operations，1 MAC = 1次加法+1次乘法操作。

因为，每次运算就等于1次FLOPs，所以，近似地，从数值上看我们有

安装thop

pip install thop

基础用法

以查看resnet50的FLOPs为例

from torchvision.models import resnet50
from thop import profile
model = resnet50()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))

查看自己模型的FLOPs

class YourModule(nn.Module):
    # your definition
def count_your_model(model, x, y):
    # your rule here
 
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ), 
                        custom_ops={YourModule: count_your_model})

提升输出结果的可读性,clever_format是更好的输出形式。
调用thop.clever_format

from thop import clever_format
flops, params = clever_format([flops, params], "%.3f")