在实验中,经常需要计算网络的参数量和预算量。
参考:https://github.com/nmhkahn/torchsummaryX
运算量和参数量的计算代码
# Calculate params & FLOPs
from torchsummaryX import summary
dummy_input = torch.zeros(1, 3, 128, 128).cuda()
summary(model.netG, dummy_input)
exit()