Pytorch打怪升级之路

仅供学习过程记录

超好用的给出模型计算量和参数量的方法

from thop import profile
input = torch.randn(1, 1, settings.DATA_SHAPE, settings.DATA_SHAPE).cuda() # 4维，B,C,H,W
flops, params = profile(net, inputs=(input,))  # 计算量和参数量
print("%s | %.2f | %.2f" % (model_name, params / (1000 ** 2), flops / (1000 ** 3))) # 单位为M

参数量的计算也可以不用import 包

print("Parameter numbers: {}".format(sum(p.numel() for p in net.parameters())/1000**2))