模型计算量 MAC/FLOPs 的手动统计方法

简介

理解神经网络中的 MAC（乘累加操作）和 FLOPs（浮点运算）对于优化网络性能和效率至关重要。通过手动计算这些指标，可以更深入地了解网络结构的计算复杂性和资源需求。这不仅能帮助设计高效的模型，还能在训练和推理阶段节省时间和资源。本文将通过实例演示如何计算全连接层(fc)、卷积层(conv) 以及 自注意力模块(self-attention) 的 FLOPs 和 MACs，并探讨其对资源效率、内存效率、能耗和模型优化的影响。

背景

为什么理解神经网络中的MAC和FLOPs很重要？

在本节中，我们将深入探讨神经网络中 MAC（乘累加操作）和 FLOPs（浮点运算）的概念。通过学习如何使用笔和纸手动计算这些指标将获得对各种网络结构的计算复杂性和效率的基本理解。

理解 MAC 和 FLOPs 不仅仅是学术练习；它是优化神经网络性能和效率的关键组成部分。它有助于设计既计算高效又有效的模型，从而在训练和推理阶段节省时间和资源。

这是一个在 Colab 笔记本中完全运行的示例

资源效率

理解 FLOPs 有助于估算神经网络的计算成本。通过优化 FLOPs 的数量，可以潜在地减少训练或运行神经网络所需的时间。

内存效率

MAC 操作通常决定了网络的内存使用情况，因为它们直接与网络中的参数和激活数量相关。减少 MACs 有助于使网络的内存使用更高效。

能耗

功耗效率

FLOPs 和 MAC 操作都对运行神经网络的硬件的功耗有贡献。通过优化这些指标，可以潜在地减少运行网络所需的能量，这对于移动设备和嵌入式设备尤为重要。

模型优化

• 剪枝和量化
  理解 FLOPs 和 MACs 可以帮助通过剪枝（去除不必要的连接）和量化（降低权重和激活的精度）等技术优化神经网络，这些技术旨在减少计算和内存成本。

性能基准

• 模型间比较
  FLOPs 和 MACs 提供了一种比较不同模型计算复杂性的方法，这可以作为为特定应用选择模型的标准。

• 硬件基准
  这些指标还可以用于对比不同硬件平台运行神经网络的性能。

• 边缘设备上的部署

  • 实时应用
    对于实时应用，特别是在计算资源有限的边缘设备上，理解和优化这些指标对于确保网络能够在应用的时间限制内运行至关重要。
  • 电池寿命
    在电池供电的设备中，减少神经网络的计算成本（从而减少能耗）可以帮助延长电池寿命。

研究与发展

• 设计新算法
  在开发新算法或神经网络结构时，研究人员可以使用这些指标作为指导，目的是在不牺牲精度的情况下提高计算效率。

FLOPs 和 MACs 定义

• FLOP（浮点运算）被认为是加法、减法、乘法或除法运算。

• MAC（乘加运算）基本上是一次乘法加上一次加法，即 MAC = a * b + c。它算作两个FLOP（一次乘法和一次加法）。

1. 全连接层 FLOPs 计算

现在，我们将创建一个包含三层的简单神经网络，并开始计算所涉及的操作。以下是计算第一层线性层（全连接层）操作数的公式：

• 对于具有 I 个输入和 O 个输出的全连接层，操作数如下：
  • MACs: I × O
  • FLOPs: 2 × (I × O)（因为每个 MAC 算作两个 FLOP）

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchprofile import profile_macs

class SimpleLinearModel(nn.Module):
    def __init__(self):
        super(SimpleLinearModel,self).__init__()
        self.fc1 = nn.Linear(in_features=10, out_features=20, bias=False)
        self.fc2 = nn.Linear(in_features=20, out_features=15, bias=False)
        self.fc3 = nn.Linear(in_features=15, out_features=1, bias=False)
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        F.relu(x)
        x = self.fc3(x)
        return x

linear_model = SimpleLinearModel().cuda()
sample_data = torch.randn(1, 10).cuda()

步骤 1：识别层参数

• 对于给定的模型，我们定义了三层线性层：
  fc1：10 个输入特征，20 个输出特征
  fc2：20 个输入特征，15 个输出特征
  fc3：15 个输入特征，1 个输出特征

步骤 2：计算 FLOPs 和 MACs

现在，计算每层的 MACs 和 FLOPs：

• 层 fc1：
  MACs = 10 × 20 = 200
  FLOPs = 2 × MACs = 2 × 200 = 400

• 层 fc2：
  MACs = 20 × 15 = 300
  FLOPs = 2 × MACs = 2 × 300 = 600

• 层 fc3：
  MACs = 15 × 1 = 15
  FLOPs = 2 × MACs = 2 × 15 = 30

步骤 3：总结结果

• 最后，为了找到单个输入通过整个网络的总 MACs 和 FLOPs，我们将所有层的结果相加：
• 总 MACs = MACs(fc1) + MACs(fc2) + MACs(fc3) = 200 + 300 + 15 = 515
• 总 FLOPs = FLOPs(fc1) + FLOPs(fc2) + FLOPs(fc3) = 400 + 600 + 30 = 1030

使用 torchprofile 库验证

可以使用 torchprofile 库来验证给定神经网络模型的 FLOPs 和 MACs 计算。以下是具体操作步骤：

macs = profile_macs(linear_model, sample_data)
print(macs)
# -> 515

2. 卷积神经网络（CNNs）

现在，让我们确定一个简单卷积模型的 MACs（乘加运算）和 FLOPs（浮点运算）。由于诸如步幅、填充和核大小等因素，这种计算比我们之前用密集层的例子更复杂一些。然而，我将逐步讲解以便于学习。

class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.fc =  nn.Linear(in_features=32*28*28, out_features=10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x

x = torch.rand(1, 1, 28, 28).cuda()
conv_model = SimpleConv().cuda()

计算卷积操作时的重要考虑因素

在计算卷积核的操作时，必须记住核的通道数量应与输入的通道数量相匹配。例如，如果我们的输入是一个有三个颜色通道的 RGB 图像，则核的维度将是 3x3x3 以匹配输入的三个通道。

为了演示的目的，我们将保持图像大小在整个卷积层中一致。为此，我们将填充和步幅值都设置为1。

第一步：确定层参数

对于给定的模型，我们定义了两个卷积层和一个线性层：

• conv1: 1 个输入通道，16 个输出通道，核大小为 3
• conv2: 16 个输入通道，32 个输出通道
• fc: 32x28x28 个输入特征，10 个输出特征。因为我们的图像在卷积层中没有改变

第二步：计算FLOPs和MACs

现在，计算每层的 MACs 和 FLOPs：

公式是：output_image_size * kernel_shape * output_channels

层conv1:

• MACs = 28 * 28 * 3 * 3 * 1 * 16 = 1,12,896
• FLOPs = 2 × MACs = 2 × 1,12,896 = 2,25,792

层conv2:

• MACs = 28 × 28 * 3 * 3 * 16 * 32 = 3,612,672
• FLOPs = 2 × MACs = 2 × 3,612,672 = 7,225,344

层fc:

• MACs = 32 * 28 * 28 * 10 = 250,880
• FLOPs = 2 × MACs = 2 × 250,880 = 501,760

第三步：汇总结果

最后，为了找到单个输入通过整个网络的总MACs和FLOPs，我们汇总所有层的结果：

• 总MACs = MACs(conv1) + MACs(conv2) + MACs(fc) = 1,12,896 + 3,612,672 + 250,880 = 3,976,448
• 总FLOPs = FLOPs(conv1) + FLOPs(conv2) + FLOPs(fc) = 2,25,792 + 7,225,344 + 501,760 = 7,952,896

使用torchprofile库验证操作

macs = profile_macs(conv_model, (x,))
print(macs)
# 输出: 3976448

3. 自注意力模块 (self-attention) FLOPs 计算

在涵盖了线性和卷积层的 MACs 之后，我们的下一步是确定自注意力模块的FLOPs（浮点运算），这是大型语言模型中的一个关键组件。这个计算对于理解这些模型的计算复杂度至关重要。让我们深入探讨。

class SimpleAttentionBlock(nn.Module):
    def __init__(self, embed_size, heads):
        super(SimpleAttentionBlock, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (
            self.head_dim * heads == embed_size
        ), "Embedding size needs to be divisible by heads"

        self.values = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.keys = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.queries = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, queries, mask):
        N = queries.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]
        print(values.shape)
        values = self.values(values).reshape(N,  self.heads, value_len, self.head_dim)
        keys = self.keys(keys).reshape(N, self.heads, key_len, self.head_dim)
        queries = self.queries(queries).reshape(N,  self.heads, query_len, self.head_dim)

        energy = torch.matmul(queries, keys.transpose(-2, -1))        

        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.nn.functional.softmax(energy, dim=3)
        out = torch.matmul(attention, values).reshape(
            N, query_len, self.heads * self.head_dim
        )

        return self.fc_out(out)

第一步：确定层参数

线性变换

让我们定义一些超参数：

batch_size = 1
seq_len = 10
embed_size = 256

在注意力块中，我们有三个线性变换（用于查询、键和值），以及一个在末尾的线性变换（fc_out）。

输入大小: [batch_size, seq_len, embed_size]

线性变换矩阵: [embed_size, embed_size]

MACs: batch_size × seq_len × embed_size × embed_size

查询、键、值线性变换：

• 查询变换的MACs = 1 × 10 × 256 × 256 = 655,360
• 键变换的MACs = 1 × 10 × 256 × 256 = 655,360
• 值变换的MACs = 1 × 10 × 256 × 256 = 655,360

能量计算: 查询（重塑后）和键（重塑后）点积——一个点积操作。

MACs: batch_size × seq_len × seq_len × heads × head_dim

查询和键的点积

MACs = 1 × 10 × 10 × 8 × 32 [32 因为256/8] = 25,600

从注意力权重和值的计算输出: 注意力权重和值（重塑后）点积——另一个点积操作。

MACs : batch_size × seq_len × seq_len × heads × head_dim

注意力和值的点积

MACs = 1 × 10 × 10 × 8 × 32 = 25,600

全连接输出（fc_out）

MACs: batch_size × seq_len × heads × head_dim × embed_size

MACs = 1 × 10 × 8 × 32 × 256 = 655,360

第二步：汇总结果

总 MACs = MACs（conv1） + MACs（conv2） + MACs（fc）= 655,360 + 655,360 + 655,360 + 25,600 + 25,600 + 655,360 = 2,672,640

总 FLOPs = 2 × 总MACs = 5,345,280

使用torchprofile库验证操作

# 创建模型实例
model = SimpleAttentionBlock(embed_size=256, heads=8).cuda()

# 生成一些样本数据（5个序列的批次，每个长度为10，嵌入大小为256）
values = torch.randn(1, 10, 256).cuda()
keys = torch.randn(1, 10, 256).cuda()
queries = torch.randn(1, 10, 256).cuda()

# 简化起见，没有掩码
mask = None

# 使用样本数据进行前向传递
macs = profile_macs(model, (values, keys, queries, mask))
print(macs)
# -> 2672640

总结：按不同批次大小缩放MACs和FLOPs

在我们的计算中，我们主要考虑了批次大小为 1。然而，按更大的批次大小缩放 MACs 和 FLOPs 是很简单的。

要计算批次大小大于 1 的 MACs 或 FLOPs，您可以简单地将批次大小 1 得到的总 MACs 或 FLOPs 乘以所需的批次大小值。此缩放允许您估计神经网络模型的各种批次大小的计算需求。

请记住，结果将直接线性缩放批次大小。例如，如果您的批次大小为 32，您可以通过将批次大小为 1 的值乘以 32 来获得 MACs 或 FLOPs。

原文链接： https://medium.com/@pashashaik/a-guide-to-hand-calculating-flops-and-macs-fa5221ce5ccc