网络参数和 flops 计算

python pytorch 计算机视觉

1. 网络参数计算

1.1 卷积层

假设输入维度为 m，输出维度为 n，卷积核大小为(k, k)

• bias=True: 一个卷积核含有的参数为$mk^{2}$，有 n 个卷积核，所以总参数为$mnk^{2}+n$
• bias=False: 没有 bias 时，不需要加上最后的 n，总参数为$mnk^{2}$

nn.Conv2d(3, 20, (3, 3), stride=1, padding=1) 其中 bias 默认为 True

上述参数为$2 \times 3^{2} \times 20 + 20 = 560$，如果设置 bias=False 就是 540

1.2 线性层

假设输入维度为 m，输出维度为 n

• bias=True: 输入为 m，输出为 n，参数为$mn+n$
• bias=False: 没有 bias 时，不需要加上最后的 n，总参数为$mn$

nn.Linear(3, 20, bias=True) bias 默认为 True

上述参数为$3 \times 20 + 20 = 80$，如果设置 bias=False 就是 60

1.3 归一化层

假设输入维度为 m

• affine=True: 有两个可学习的参数$\alpha$和$\beta$，参数为 2m
• affine=False: 没有参数，0

nn.BatchNorm2d(20, affine=True) affine 默认为 True

上述参数为$2 \times 20 = 40$，如果 affine=False 为 0

2. Flops

Flops 的运算为：一次乘加运算为一次 Flops

2.1 卷积层

假设输入维度为 m，输出维度为 n，卷积核大小为(k, k)，输出特征图的尺寸为(H，W)

一个卷积核每进行一次卷积所需要的乘法是$mk^{{2}$，加法为$m(k}{2}-1)+(m-1)=mk^{{2}-1$，所以一个卷积核的计算量为$mk}{2}$（忽略掉那-1 的运算，乘加就相等了），输出的特征图为 HW，输出通道为 n，所以总共的 Flops 为：$HWn \times mk^{2}$

2.2 线性层

假设输入维度为 m，输出维度为 n，batch 为 b

从下图可以看出 b=128，我们先不管这个参数，他的每一个数据 m=20，n=30，b 中的每一个数据都需要计算 m 次乘法，计算完之后如果有 bias 需要额外计算 m 次加法，输出维度为 n，所以一共需要计算 mn 次乘法和 mn 次加法，即 Flops：$mn$，算上 batch 为$bmn$

2.3 归一化

假设输出维度为 m，Batch 为 b

归一化层有可学习的参数$\alpha$和$\beta$，各有 m 次乘加操作，所以 Flops：$m$，算上 batch 为$bm$

3. 计算的模块

3.1 torchsummary 模块

pip install torchsummary

from torchsummary import summary
summary(model, (3, 224, 224))

3.2 flops_counter 模块

GitHub 地址：https://github.com/sovrasov/flops-counter.pytorch

pip install ptflops

from ptflops import get_model_complexity_info
flops, params = get_model_complexity_info(model, (3, 224, 224), as_strings=True, print_per_layer_stat=True)

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Dec 19, 2024 12:13 pm