OneFlow’s FLOPs
2023-1-19
| 2023-2-16
0  |  阅读时长 0 分钟
type
status
date
slug
summary
tags
icon
password

flowflops: OneFlow 模型的 Flops 计算

用于计算 OneFlow 模型的 FLOPs 和 Parameters 的第三方库。
源码地址: https://github.com/Oneflow-Inc/flow-OpCounter

介绍 & 使用

FLOPs

有许多人分不清楚 FLOPs 和 MACs 之间的关系,如ptflops中的issue
针对该问题,可以查看thop中的解释,翻译如下:

MACs, FLOPs, what is the difference?

FLOPs浮动算子(floating operations)的缩写,包括mul / add / div …等。
MACs 代表执行的乘法累加运算,例如: a <- a + (b x c)
如文中所示,一MACs有一mul和一add。这就是为什么在许多地方FLOPs几乎是两倍MACs的原因。
然而,现实世界中的应用要复杂得多。让我们考虑一个矩阵乘法示例。A是一个形状为 m × n 的矩阵,B是一个 n × 1 的向量。
它会是m*n个MACs和2m*n个FLOPs。但是这样的矩阵乘法实现速度很慢,需要并行化才能运行得更快。
此时MACs数值不再是 m*n 。
在比较 MAC / FLOP 时,我们希望数字与实现无关并且尽可能通用。因此在 (thop)[https://github.com/Lyken17/pytorch-OpCounter] 中,我们只考虑乘法的次数,而忽略所有其他操作。

安装方法

使用方法

目前支持两种 FLOPs 计算策略: 在 Eager 模式下计算和在 Graph 模式下计算。
在 Graph 模式下计算耗时较长,但结果更加精确
示例:
输出:
可以看到两种计算方式下的输出有一定差别,这是因为在 ResNet 的 forward 代码里存在类似 out += identity 的语句,这会造成 FLOPs 额外增加。而在 Eager 模式下我们只关注在 __init__() 中定义的网络层,所以这种情况不会在 Eager 模式中被 hook 到。
我们可以计算一下有哪些 add_n 算子在 Eager 模式中被我们忽略了:
一共为 5,519,360 ,刚好为两种模式的输出差值 4127444456 - 4121925096 = 5519360
在 Eager 模式下也会存在一些小误差,一般认为 ResNet50 的 FLOPs 为 4.09G ,而这里计算得到 4.12G ,是因为一般研究中会忽略类似 ReLU 等算子的 FLOPs 计算,所以与真实数值会有一定误差。有关一般都忽略了哪些算子的计算,可以查看 fvcore 的输出。
在ptflops包中也存在这样的问题,笔者也有在issue中回复,详见issue: https://github.com/sovrasov/flops-counter.pytorch/issues/94

Eager & Graph 模式下的 Flops 计算

接下来我们以简单修改后的 ResNet18 中的 BasicBlock 为例介绍一下两种 FLOPs 计算方式,设定网络如下:
我们统一假定输入形状为(1, 32, 64, 64)

Eager

在 Eager 模式中,我们只关注 __init__() 中定义的网络层,也就是

二维卷积

卷积的原理在此不再赘述,直接给出计算公式: 2k2 × Hout × Wout × Cin × Cout ÷ Groups

归一化

batchnorm 主要计算了均值、方差,并基于此对特征进行归一化与仿射变换,其 FLOPs 为 2 × C × H × W
如果不进行仿射变换,则其 FLOPs 为 C × H × W

激活函数

relu 对输入(1, C, H, W)进行了 y = x if x > 0 else 0 操作,也就是其 FLOPs 为 C × H × W

线性层

线性层输入为 (1, C, H, W)
线性层权重为 (W1, W)
两者相乘的 FLOPs 为 C × H × W1 × W
其本质与 matmul 计算相当

Graph

在 Graph 模式中,我们会将 flow.nn.Module 编译为 flow.nn.Graph ,从 Graph 中抽取出每一个算子输入的张量形状后再对网络的 FLOPs 进行计算
上述网络转换后的 Graph:
Graph 中由 OPERATOR 开始的层就是我们所需要的信息

卷积

flow.nn.Graphconv3x3conv1x1 会被拆解为 conv2d + bias_add(if bias==True)
由于我们只关注的卷积层的输入,而在计算 FLOPs 时需要得到卷积层输出的特征尺寸,所以我们需要依据输入计算输出特征的分辨率,方法如下
随后即可正常计算 FLOPs
至于为什么不直接得到算子输出的形状,因为解析 Graph 需要占用更多的额外时间

归一化

flow.nn.Graphnorm_layer(bn) 是一个单独的算子,其计算方法与 Eager 模式中保持一致
需要注意的是 InstanceNorm 和 GroupNorm 在 flow.nn.Graph 中将被拆解为若干胶水算子,需要逐个计算

激活函数

flow.nn.Graphrelu 是一个单独的算子,其 FLOPs 计算方法与 Eager 模式中保持一致

线性层

flow.nn.Graphlinear 会被拆解为 matmul + broadcast_add(if bias==True),其 FLOPs 计算公式与 Eager 模式中基本一致

目前支持的 Op 与模型

目前该工具支持绝大部分算子、网络层与大多数 CNN ,列表如下

Eager

Graph

FlowVision 中部分模型的计算结果

总结

简单介绍 OneFlow 模型中如何计算网络的 FLOPs
  • AI System
  • Unsupervised(unpaired) learning in blind super resolutionDebug about fused shift-window
    目录