PyTorch thop包计算模型Params和FLOPs

安装（推荐第二个方法）：

pip install thop
pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git

1.基本用法：

from torchvision.models import resnet50 
from thop import profile 

model = resnet50() 
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs = (input, ))

2.定义第三方模块的规则： 

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})

3.提高输出可读性：

调用thop.clever_format以提供更好的输出格式

from thop import clever_format

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

PyTorch parameters()也可以计算参数量：

from thop import profile
from utils.ShuffleNetV2 import shufflenetv2 
import torch 
  
model_shuffle = shufflenetv2(width_mult=0.5)
model = torch.nn.DataParallel(model_shuffle)
total = sum([param.nelement() for param in model.parameters()])
print("Number of parameter: %.2fM" % (total / 1e6))