两行代码统计模型参数量与FLOPs,这个PyTorch小工具值得一试

你的模型到底有多少参数,每秒的浮点运算到底有多少,这些你都知道吗?近日,GitHub 开源了一个小工具,它可以统计 PyTorch 模型的参数量与每秒浮点运算数(FLOPs)。有了这两种信息,模型大小控制也就更合理了。

其实模型的参数量好算,但浮点运算数并不好确定,我们一般也就根据参数量直接估计计算量了。但是像卷积之类的运算,它的参数量比较小,但是运算量非常大,它是一种计算密集型的操作。反观全连接结构,它的参数量非常多,但运算量并没有显得那么大。

此外,机器学习还有很多结构没有参数但存在计算,例如最大池化Dropout 等。因此,PyTorch-OpCounter 这种能直接统计 FLOPs 的工具还是非常有吸引力的。

OpCouter

PyTorch-OpCounter 的安装和使用都非常简单,并且还能定制化统计规则,因此那些特殊的运算也能自定义地统计进去。

我们可以使用 pip 简单地完成安装:pip install thop。不过 GitHub 上的代码总是最新的,因此也可以从 GitHub 上的脚本安装。

对于 torchvision 中自带的模型,Flops 统计通过以下几行代码就能完成:

from torchvision.models import resnet50from thop import profilemodel = resnet50()input = torch.randn(1, 3, 224, 224)flops, params = profile(model, inputs=(input, ))

我们测试了一下 DenseNet-121,用 OpCouter 统计了参数量与运算量。API 的输出如下所示,它会告诉我们具体统计了哪些结构,它们的配置又是什么样的。

最后输出的浮点运算数和参数量分别为如下所示,换算一下就能知道 DenseNet-121 的参数量约有 798 万,计算量约有 2.91 GFLOPs。

flops: 2914598912.0parameters: 7978856.0

OpCouter 是怎么算的

我们可能会疑惑,OpCouter 到底是怎么统计的浮点运算数。其实它的统计代码在项目中也非常可读,从代码上看,目前该工具主要统计了视觉方面的运算,包括各种卷积、激活函数池化批归一化等。例如最常见的二维卷积运算,它的统计代码如下所示:

def count_conv2d(m, x, y):    x = x[0]    cin = m.in_channels    cout = m.out_channels    kh, kw = m.kernel_size    batch_size = x.size()[0]    out_h = y.size(2)    out_w = y.size(3)    # ops per output element    # kernel_mul = kh * kw * cin    # kernel_add = kh * kw * cin - 1    kernel_ops = multiply_adds * kh * kw    bias_ops = 1 if m.bias is not None else 0    ops_per_element = kernel_ops + bias_ops    # total ops    # num_out_elements = y.numel()    output_elements = batch_size * out_w * out_h * cout    total_ops = output_elements * ops_per_element * cin // m.groups    m.total_ops = torch.Tensor([int(total_ops)])

总体而言,模型会计算每一个卷积核发生的乘加运算数,再推广到整个卷积层级的总乘加运算数。

定制你的运算统计

有一些运算统计还没有加进去,如果我们知道该怎样算,那么就可以写个自定义函数。

class YourModule(nn.Module):    # your definitiondef count_your_model(model, x, y):    # your rule hereinput = torch.randn(1, 3, 224, 224)flops, params = profile(model, inputs=(input, ),                        custom_ops={YourModule: count_your_model})

最后,作者利用这个工具统计了各种流行视觉模型的参数量与 FLOPs 量:

工程PyTorchGitHub
相关数据
池化技术

池化(Pooling)是卷积神经网络中的一个重要的概念,它实际上是一种形式的降采样。有多种不同形式的非线性池化函数,而其中“最大池化(Max pooling)”是最为常见的。它是将输入的图像划分为若干个矩形区域,对每个子区域输出最大值。直觉上,这种机制能够有效的原因在于,在发现一个特征之后,它的精确位置远不及它和其他特征的相对位置的关系重要。池化层会不断地减小数据的空间大小,因此参数的数量和计算量也会下降,这在一定程度上也控制了过拟合。通常来说,CNN的卷积层之间都会周期性地插入池化层。

激活函数技术

在 计算网络中, 一个节点的激活函数定义了该节点在给定的输入或输入的集合下的输出。标准的计算机芯片电路可以看作是根据输入得到"开"(1)或"关"(0)输出的数字网络激活函数。这与神经网络中的线性感知机的行为类似。 一种函数(例如 ReLU 或 S 型函数),用于对上一层的所有输入求加权和,然后生成一个输出值(通常为非线性值),并将其传递给下一层。

机器学习技术

机器学习是人工智能的一个分支,是一门多领域交叉学科,涉及概率论、统计学、逼近论、凸分析、计算复杂性理论等多门学科。机器学习理论主要是设计和分析一些让计算机可以自动“学习”的算法。因为学习算法中涉及了大量的统计学理论,机器学习与推断统计学联系尤为密切,也被称为统计学习理论。算法设计方面,机器学习理论关注可以实现的,行之有效的学习算法。

Dropout技术

神经网络训练中防止过拟合的一种技术

最大池化技术

最大池化(max-pooling)即取局部接受域中值最大的点。

批归一化技术

批归一化(Batch Normalization,BN)由谷歌于2015年提出,是一个深度神经网络训练的技巧,它不仅可以加快模型的收敛速度,还能在一定程度上缓解深层网络中的“梯度弥散”问题,从而使得训练深层网络模型更加容易和稳定。目前BN已经成为几乎所有卷积神经网络的标配技巧了。从字面意思看来Batch Normalization(简称BN)就是对每一批数据进行归一化。

推荐文章
暂无评论
暂无评论~