6.1.3.4. 计算量工具¶
6.1.3.4.1. 计算量的定义¶
计算量是用来评估神经网络大小的常用工具。在常见的计算量工具中一般只统计两种操作的计算量,一个是卷积相关,另一个是全连接相关的。
以常见的 torch.nn.Conv2d
为例,输入数据的形状为 bs * c_in * w_in * h_in
,输出数据的形状为bs * c_out * w_out * h_out
,卷积核的大小为 f
。则该 Conv2d
的计算量为 2 * bs * f * f * c_in * c_out * w_out * h_out
。 2
表示加法计算量和乘法计算量各一半。
而 torch.nn.Linear
的情况里,输入数据的神经元个数为 c_in
,输出数据的神经元个数为 c_out
。其实这种全连接层可以作为一个特殊的卷积层,输入输出的大小均为 1x1
,卷积核大小也是 1x1
。则全连接层的计算量为 bs * (c_in * c_out + c_out)
,需要注意是这里乘法和加法的计算量并不完全一致。
注意
量化模型和QAT模型和对应的浮点模型的计算量是完全一致的。
6.1.3.4.2. 计算量工具的使用方法¶
目前计算量工具支持统计计算量的op
就是上面提到的两种op
,分别为torch.nn.Conv2d
和torch.nn.Linear
。
python3 tools/calops.py --config ${CONFIG_PATH}
其中 config
中影响计算量的主要key是 test_model
(或者 model
),以及 test_inputs
。 model
相关的决定了计算量工具需要检查的模型,而test_inputs
决定了输入的大小,除此之外,输入形状还可以通过--input-shape B,C,H,W
的输入参数决定。