理解神经网络中的 MAC(乘累加操作)和 FLOPs(浮点运算)对于优化网络性能和效率至关重要。通过手动计算这些指标,可以更深入地了解网络结构的计算复杂性和资源需求。这不仅能帮助设计高效的模型,还能在训练和推理阶段节省时间和资源。本文将通过实例演示如何计算全连接层(fc)、卷积层(conv) 以及 自注意力模块(self-attention) 的 FLOPs 和 MACs,并探讨其对资源效率、内存效率、能耗和模型优化的影响。
在本节中,我们将深入探讨神经网络中 MAC(乘累加操作)和 FLOPs(浮点运算)的概念。通过学习如何使用笔和纸手动计算这些指标将获得对各种网络结构的计算复杂性和效率的基本理解。
理解 MAC 和 FLOPs 不仅仅是学术练习;它是优化神经网络性能和效率的关键组成部分。它有助于设计既计算高效又有效的模型,从而在训练和推理阶段节省时间和资源。
这是一个在 Colab 笔记本中完全运行的示例
理解 FLOPs 有助于估算神经网络的计算成本。通过优化 FLOPs 的数量,可以潜在地减少训练或运行神经网络所需的时间。
MAC 操作通常决定了网络的内存使用情况,因为它们直接与网络中的参数和激活数量相关。减少 MACs 有助于使网络的内存使用更高效。
FLOPs 和 MAC 操作都对运行神经网络的硬件的功耗有贡献。通过优化这些指标,可以潜在地减少运行网络所需的能量,这对于移动设备和嵌入式设备尤为重要。
模型间比较
FLOPs 和 MACs 提供了一种比较不同模型计算复杂性的方法,这可以作为为特定应用选择模型的标准。
硬件基准
这些指标还可以用于对比不同硬件平台运行神经网络的性能。
边缘设备上的部署
FLOP(浮点运算)被认为是加法、减法、乘法或除法运算。
MAC(乘加运算)基本上是一次乘法加上一次加法,即 MAC = a * b + c。它算作两个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()
现在,计算每层的 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
可以使用 torchprofile 库来验证给定神经网络模型的 FLOPs 和 MACs 计算。以下是具体操作步骤:
macs = profile_macs(linear_model, sample_data) print(macs) # -> 515
现在,让我们确定一个简单卷积模型的 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。
对于给定的模型,我们定义了两个卷积层和一个线性层:
现在,计算每层的 MACs 和 FLOPs:
公式是:output_image_size * kernel_shape * output_channels
层conv1:
层conv2:
层fc:
最后,为了找到单个输入通过整个网络的总MACs和FLOPs,我们汇总所有层的结果:
macs = profile_macs(conv_model, (x,)) print(macs) # 输出: 3976448
在涵盖了线性和卷积层的 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
查询、键、值线性变换:
1 × 10 × 256 × 256 = 655,360
1 × 10 × 256 × 256 = 655,360
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
# 创建模型实例 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
在我们的计算中,我们主要考虑了批次大小为 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