1、导入所需要的包
import ...
from thop import profile
2、定义模型
net = UNet(3, 2)
3、输入大小
// 与输入大小相同,第一维为batch size
input = torch.randn(1, 3, 256, 256)
4、计算计算量和参数量
// inputs与模型输入保持一致
Flops, params = profile(net, inputs=(input,))
// 模型一次输入两证图片如下所示(模型的输入为两张图片,并不是batch size为2)
Flops, params = profile(net, inputs=(input,input))
5、输出计算量和参数量
print('Flops: % .4fG'%(Flops / 1000000000))# 计算量
print('params参数量: % .4fM'% (params / 1000000)) #参数量
6、可能存在的问题:
报错:ModuleNotFoundError: No module named ‘thop’
解决:
pip install thop