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

it2022-05-05  158

神经网络参数量统计工具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时间作为加速效果的衡量指标。


最新回复(0)