Pytorch计算模型参数总量和模型计算量

最新推荐文章于 2025-05-30 11:00:49 发布

confusingbird

最新推荐文章于 2025-05-30 11:00:49 发布

阅读量1.8w

点赞数 23

CC 4.0 BY-SA版权

文章标签： Pytorch

本文链接：https://blog.youkuaiyun.com/confusingbird/article/details/103914102

计算这两个值之前要进行模型加载，模型加载的方法见https://blog.youkuaiyun.com/confusingbird/article/details/103913915，

加载完模型后得到模型model

方法1：pytorch自带方法，计算模型参数总量

 total = sum([param.nelement() for param in model.parameters()])

 print("Number of parameter: %.2fM" % (total/1e6))

方法2：FeatherNet作者提供的一个方法，计算模型参数总量和模型计算量

def count_params(model, input_size=224):
    # param_sum = 0
    with open('models.txt', 'w') as fm:
        fm.write(str(model))

    # 计算模型的计算量
    calc_flops(model, input_size)

    # 计算模型的参数总量
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])

    print('The network has {} params.'.format(params))


# 计算模型的计算量
def calc_flops(model, input_size):

    def conv_hook(self, input, output):
        batch_size, input_channels, input