Pytorch中统计模型的计算量(FLOPs)和参数量(Params)

1.计算量和参数量

1.1 计算量和参数量的含义

神经网络的计算量和参数量是评估其性能和效率的两个重要指标。计算量和参数量分别对应时间复杂度和空间复杂度。
参数量： 参数量指的是神经网络中所有可学习的权重和偏置的总和。在训练过程中，这些参数会通过梯度下降等优化算法进行调整。参数量直接关系到模型的大小，以及模型对计算资源的占用。一般来说，参数量越大，模型的表达能力越强，但同时也意味着更高的计算和存储需求，更容易发生过拟合。
计算量： 计算量主要是指在神经网络的前向传播和反向传播过程中所需的浮点运算次数。这包括权重与输入的矩阵乘法、激活函数的计算等。计算量决定了模型运行的速度，对于实时应用来说尤其重要。计算量大的模型需要更长的训练时间，也要求更强的计算能力。

1.2 计算量的单位

在讨论神经网络的计算量时，常用的单位有MAC (Multiply-Accumulate Operations) 和 FLOPs (Floating Point Operations)。这两个单位用于衡量执行计算操作的数量，但它们之间有一些区别：

MAC (Multiply-Accumulate Operations): MAC 是指一个乘法操作后紧跟一个加法操作的数量。在深度学习中，MAC 通常用于衡量卷积操作的计算量，因为卷积操作涉及到权重和输入特征的乘法，然后将结果累加起来。
FLOPs (Floating Point Operations): FLOPs 是指浮点运算的总数。这个术语通常用于更广泛地描述任何涉及浮点数的计算操作。在深度学习中，FLOPs 用于衡量模型的总计算量，包括所有的乘法、加法和其他可能的操作（如激活函数的计算）。
由于每个 MAC 包含一个乘法和一个加法，所以 FLOPs 将是 MAC 计算量的两倍。

1.3 参数量的单位

参数量。每个参数可以是权重（在神经网络中的连接上）或者偏置项。参数量直接关系到模型的大小，以及模型对计算资源和存储的需求。

在深度学习中，参数量也常常用不同的单位来表示，以适应不同规模的模型：
千（K）：1K = 1,000参数
百万（M）：1M = 1,000,000参数
十亿（B）：1B = 1,000,000,000参数

多层感知机 层数较少的时候，以K作为单位。
深度神经网络常以M作为单位。
大语言模型 如chatglm-6B 以B作为单位。

2.使用库函数统计神经网络的运算量和参数量

2.1 thop库

单输入

import torch
from torchsummary import summary
from thop import profile
from thop import clever_format
import torchvision.models as models

# 假设我们有一个预训练的模型
model = models.resnet50()
model.eval()

# 使用thop分析模型的运算量和参数量
input = torch.randn(1, 3, 224, 224)  # 随机生成一个输入张量，这个尺寸应该与模型输入的尺寸相匹配
flops, params = profile(model, inputs=(input,))

# 将结果转换为更易于阅读的格式
flops, params = clever_format([flops, params], '%.3f')

print(f"运算量：{flops}, 参数量：{params}")

在上面的代码中，我们首先导入了必要的库，并加载了一个预训练的ResNet-18模型。然后，我们创建了一个随机张量作为模型的输入，并使用profile函数来计算模型的运算量（FLOPs）和参数量。最后，我们使用clever_format函数将结果转换为更易于阅读的格式，并将其打印出来。

请注意，为了正确地分析模型，你需要确保输入张量的尺寸与模型期望的输入尺寸相匹配。如果你要分析的是自己的模型，那么你需要根据模型的实际情况来调整输入张量的尺寸。

补充模型输入有多个的情况

import torch
from thop import profile
from thop import clever_format

def conv_output_size(input_size, kernel_size, padding, stride):
    return (input_size - kernel_size + 2 * padding) // stride + 1

# 定义一个多输入的模型
class MultiInputModelFixedAgain(torch.nn.Module):
    def __init__(self):
        super(MultiInputModelFixedAgain, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1, stride=1)
        self.conv2 = torch.nn.Conv2d(3, 64, kernel_size=3, padding=1, stride=1)
        # 计算每个卷积层的输出特征图大小
        input_size = 224
        kernel_size = 3
        padding = 1
        stride = 1
        output_size = conv_output_size(input_size, kernel_size, padding, stride)
        features_per_conv = 64 * output_size ** 2
        total_features = 2 * features_per_conv
        self.fc = torch.nn.Linear(total_features, 10)

    def forward(self, x1, x2):
        x1 = self.conv1(x1)
        x2 = self.conv2(x2)
        x = torch.cat((x1, x2), dim=1)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 创建模型实例
model_fixed_again = MultiInputModelFixedAgain()

# 创建两个随机输入张量
input1 = torch.randn(1, 3, 224, 224)
input2 = torch.randn(1, 3, 224, 224)

# 使用thop分析模型的运算量和参数量
flops, params = profile(model_fixed_again, inputs=(input1, input2))

# 将结果转换为更易于阅读的格式
flops, params = clever_format([flops, params], '%.3f')

print(f"运算量：{flops}, 参数量：{params}")

2.2 ptflops 库

下面是如何使用 ptflops 来计算模型 FLOPs 的示例：

from ptflops import get_model_complexity_info
import torch
import torchvision.models as models

# 创建一个模型实例，这里以 ResNet-50 为例
model = models.resnet50()

# 输入模型的尺寸，这个尺寸应该与模型输入的尺寸相匹配
input_res = (3, 224, 224)

# 使用 ptflops 计算模型的 FLOPs 和参数量
macs, params = get_model_complexity_info(model, input_res, as_strings=True, print_per_layer_stat=True)

print(f"模型 FLOPs: {macs}")
print(f"模型参数量: {params}")

在上面的代码中，我们首先从 torchvision.models 导入了一个预训练的 ResNet-50 模型。然后，我们定义了模型的输入尺寸，并使用 get_model_complexity_info 函数来计算模型的乘积累加操作（MACs，可以近似为 FLOPs）和参数量。最后，我们将结果打印出来。

get_model_complexity_info 函数的参数 as_strings=True 表示返回的结果以字符串形式表示，print_per_layer_stat=True 表示打印每个层的统计信息。如果你不想打印每个层的详细信息，可以将 print_per_layer_stat 设置为 False

2.3 pytorch_model_summary库

from pytorch_model_summary import summary
import torch
import torchvision.models as models

# 创建一个模型实例，这里以 ResNet-50 为例
model = models.resnet50()

# 定义一个假的输入张量，这个尺寸应该与模型输入的尺寸相匹配
input_tensor = torch.rand(1, 3, 224, 224)

# 使用 pytorch_model_summary 获取模型的摘要
model_summary = summary(model, input_tensor)

print(model_summary)

在上面的代码中，我们首先从 torchvision.models 导入了一个预训练的 ResNet-50 模型。然后，我们定义了一个假的输入张量，并使用 summary 函数来获取模型的摘要。col_names 参数用于指定要在摘要中显示的列，verbose 参数用于控制输出的详细程度。
pytorch_model_summary 库提供了一个简洁的表格形式的输出，其中包含了每层的输入尺寸、输出尺寸、参数量、卷积核尺寸和 MACs。这有助于开发者快速了解模型的结构和性能。
pytorch_model_summary 库可能不支持所有类型的模型或层。如果在使用过程中遇到问题，可能需要查看库的文档或源代码，以了解它支持哪些模型和层。