模型压缩之神经网络模型参数量统计工具thop

神经网络参数量统计工具thop

安装方式：

pip install thop

使用方法，话不多说，直接上代码。

import torch
from torchvision import models
from thop import profile

model = models.densenet121()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))

运行的最终结果为模型的计算量和参数量——评价模型压缩方法的重要指标。

Flops of DenseNet-121 is  2913996800.0
Parameters of DenseNet-121 is  7978856.0

当然，也可以添加自己定义的网络模块，前提是你知道参数量和计算量计算的方法。

class YourModule(nn.Module):
    # your definition
def count_your_model(model, x, y):
    # your rule here

input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ), 
                        custom_ops={YourModule: count_your_model})

官方结果

PyTorch-OpCounter GitHub 地址：https://github.com/Lyken17/pytorch-OpCounter

模型压缩评价指标

模型压缩将会是大数据任务（Big Data）的一种常用手段，所要解决的问题就是庞大的模型与有限的计算资源之间的矛盾。官方措辞“在过去的几年中，我们见证了（卷积）神经网络在xxx领域所取得的成功。”精心设计的网络结构功不可没，但很大程度上还是得益于网络的加深、加大。这就导致了网络的参数量暴增，为训练和应用带来挑战。模型训练可以使用更多的GPU（$$$↑↑），但是在应用场景下就不一样了，受制于成本、体积、功耗等因素，有时候应用平台的硬件性能较差（算力低，带宽小，内存小），偏偏任务的实时性要求很高（如人脸识别）。既然硬件资源受限，那么必然从模型的角度入手，减少模型参数量和计算量。

目前，模型压缩的主流方法主要分为模型剪枝（权重剪枝/粗粒度，滤波器剪枝/细粒度）、结构设计（深度可分离卷积MobileNet，通道混排ShuffleNet等）以及量化。（具体综述可以看这篇文章，总结的不错https://arxiv.org/abs/1904.00346?context=cs.CV）

参数量和计算量只是理论上的压缩比与加速比，具体的加速效果还要结合特定的硬件平台。值得注意的是，已经有模型压缩方向相关的论文开始以CPU时间作为加速效果的衡量指标。