获取网络模型的每一层参数量与计算量（Flops）———Pytorch

曾经终败给现在 2023-07-18 06:22 26阅读 0赞

## 一、前言 ##

在现在AI各种技术需要luo地的时期，网络模型大小能够满足嵌入式平台极为重要，不仅仅需要看模型效果，也要看模型的计算量与参数，所以在评估模型的时候就要分析网络的参数量与计算量；

## 二、推荐pytorch工具 ##

###       1、ptflops ###

安装： pip3 install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git

使用：具体使用参照 [https://github.com/sovrasov/flops-counter.pytorch][https_github.com_sovrasov_flops-counter.pytorch]，在第三节我也贴上了我使用代码，有具体解释与注意事项，可看后面

评价：博主就是用的这个，他可以直接看每一层的参数量与计算量，每一层的参数计算量占模型所有参数的百分比，且有每一层的卷积步长、核大小、输入输出通道数量等，并且也会输出总量；推荐

###      2、thop ###

安装：pip3 install thop    或者  pip3 install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git

使用：同样也可参照官方：[https://github.com/Lyken17/pytorch-OpCounter][https_github.com_Lyken17_pytorch-OpCounter]，也可以看我使用的代码

评价：也是很多人用的，可以自定义网络层计算，可看官方使用方法，也可以输出参数的总量与计算量，但是博主没发现如何输出每一层的计算量参数量，但是看到了借口，改源码可能需要编译，但是没时间就没弄了，以后有时间再研究；

## 三、使用实例 ##

###      1、ptflops使用 ###

from torchvision.models import resnet50
    import torch
    import torchvision.models as models
    # import torch
    from ptflops import get_model_complexity_info
    
    # model = models.resnet50()   #调用官方的模型，
    checkpoints = '自己模型的path'
    model = torch.load(checkpoints)
    model_name = 'yolov3 cut'
    flops, params = get_model_complexity_info(model, (3,320,320),as_strings=True,print_per_layer_stat=True)
    print("%s |%s |%s" % (model_name,flops,params))

**\#注意，这里输入一定是要tuple类型，且不需要输入batch，直接输入输入通道数量与尺寸，如（3,320,320）  320为网络输入尺寸；print\_per\_layer\_stat这个参数表示是否打印每一层的参数量与计算量；输出为网络模型的总参数量（单位M，即百万）与计算量（单位G，即十亿）**

输出样例；

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2NvbXdheV9MaQ_size_16_color_FFFFFF_t_70][]

### 2、thop使用 ###

from torchvision.models import resnet50
    from thop import profile
    
    # model = resnet50()
    checkpoints = '模型path'
    model = torch.load(checkpoints)
    model_name = 'yolov3 cut asff'
    input = torch.randn(1, 3, 224, 224)
    flops, params = profile(model, inputs=(input, ),verbose=True)
    print("%s | %.2f | %.2f" % (model_name, params / (1000 ** 2), flops / (1000 ** 3)))#这里除以1000的平方，是为了化成M的单位，

**注意:输入必须是四维的；profile函数中verbose参数，设置为true，我认为就是输出每一层的参数量与计算量的，但是没输出，下方代码为profile函数的源码，可看到31行verbose参数，待研究；**

def profile(model, inputs, custom_ops=None, verbose=True):
        handler_collection = []
        if custom_ops is None:
            custom_ops = {}
    
        def add_hooks(m):
            if len(list(m.children())) > 0:
                return
    
            if hasattr(m, "total_ops") or hasattr(m, "total_params"):
                logger.warning("Either .total_ops or .total_params is already defined in %s. "
                               "Be careful, it might change your code's behavior." % str(m))
    
            m.register_buffer('total_ops', torch.zeros(1))
            m.register_buffer('total_params', torch.zeros(1))
    
            for p in m.parameters():
                m.total_params += torch.Tensor([p.numel()])
    
            m_type = type(m)
            fn = None
            if m_type in custom_ops:  # if defined both op maps, use custom_ops to overwrite.
                fn = custom_ops[m_type]
            elif m_type in register_hooks:
                fn = register_hooks[m_type]
    
            if fn is None:
                if verbose:
                    logger.info("THOP has not implemented counting method for ", m)
            else:
                if verbose:
                    logger.info("Register FLOP counter for module %s" % str(m))
                handler = m.register_forward_hook(fn)
                handler_collection.append(handler)
    
        training = model.training
    
        model.eval()
        model.apply(add_hooks)
    
        with torch.no_grad():
            model(*inputs)
    
        total_ops = 0
        total_params = 0
        for m in model.modules():
            if len(list(m.children())) > 0:  # skip for non-leaf module
                continue
            total_ops += m.total_ops
            total_params += m.total_params
    
        total_ops = total_ops.item()
        total_params = total_params.item()
    
        # reset model to original status
        model.train(training)
        for handler in handler_collection:
            handler.remove()
    
        # remove temporal buffers
        for n, m in model.named_modules():
            if len(list(m.children())) > 0:
                continue
            if "total_ops" in m._buffers:
                m._buffers.pop("total_ops")
            if "total_params" in m._buffers:
                m._buffers.pop("total_params")
    
        return total_ops, total_params

**提高输出可读性**

加入一下代码

from thop import clever_format
    macs, params = clever_format([flops, params], "%.3f")

[https_github.com_sovrasov_flops-counter.pytorch]: https://github.com/sovrasov/flops-counter.pytorch
[https_github.com_Lyken17_pytorch-OpCounter]: https://github.com/Lyken17/pytorch-OpCounter
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2NvbXdheV9MaQ_size_16_color_FFFFFF_t_70]: /images/20230528/0f840963d7844510b6407c634deda99c.png