计算模型参数量

参考： https://blog.csdn.net/weixin_43379058/article/details/108433197

 

tensorflow

        model = CPASSRnet(sess, args)
        num_params  = 0
        for variable in tf.trainable_variables():
            shape = variable.get_shape()
            num_params += reduce(mul, [dim.value for dim in shape], 1)
            print(variable)
        print('num_params', num_params / 1024 / 1024) 

tensorflow

        model = CPASSRnet(sess, args)
        total = np.sum([np.prod(v.get_shape().as_list()) for v in tf.trainable_variables()])
        print('+ Number of params: %.2fM' % (total / 1024/1024))

pytorch

cfg = parse_args()
net = PASSRnet(cfg.scale_factor).to(cfg.device)
total_params = sum(p.numel() for p in net.parameters())
print(f'{total_params/1024/1024:,} total parameters.')