利用thop计算模型的参数量params和计算量flops

from thop import profile

s = b, c, d, h, w = 1, 3, 5, 224, 224
t = torch.ones(s).cuda()

net = CoarseNet().cuda()
output, spatial_domain_feature = net(t)
F_net = F_net().cuda()
out_f, temporal_domain_feature = F_net(output, t)

flops, params = profile(net,inputs=(t,))
flops1, params1 = profile(F_net,inputs=(output, t,))
print('flops: %.2f M, params: %.2f M' % (flops / 1000000.0, params / 1000000.0))
print('flops: %.2f M, params: %.2f M' % (flops1 / 1000000.0, params1 / 1000000.0))

利用thop计算模型的参数量params和计算量flops。

有多个输入时，只需要将input(output,)更改为input(output1,output2,)即可。非常方便。