【pytorch】评估计算 PyTorch 模型的 FLOPs 和参数量

在深度学习模型中，评估模型的复杂性和效率时，FLOPs（每秒浮点运算次数）和参数量是两个重要指标。FLOPs 可以衡量模型的计算量，参数量则决定模型的规模和存储需求。

将以一个简单的卷积神经网络（CNN）为例，演示如何计算这些指标。

模型结构

这是一个简单的 CNN 模型，包含两个卷积层、ReLU 激活函数、Batch Normalization 层和一个全连接层：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(128 * 32 * 32, 1000)  # 全连接层

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = x.view(x.size(0), -1)  # 将卷积输出展平成2D张量
        x = self.fc(x)
        return x

方法1: 使用 thop 计算 FLOPs 和参数量

thop 是一个方便的库，可以用来计算 PyTorch 模型的 FLOPs 和参数量。下面展示了如何使用 thop 来计算这些指标：

安装 thop

你可以使用以下命令安装 thop：

pip install thop

代码示例

import torch
from thop import profile
from torch import nn

# 定义模型
model = MyModel()

# 创建一个输入张量
input = torch.randn(1, 3, 32, 32)  # 假设输入尺寸为 (32, 32)

# 计算 FLOPs 和参数量
flops, params = profile(model, inputs=(input,))

print(f"FLOPs: {flops}")
print(f"Parameters: {params}")

输出结果示例

FLOPs: 402341824
Parameters: 131148032

方法2: 使用 ptflops 计算 FLOPs 和参数量

ptflops 是另一个计算 PyTorch 模型 FLOPs 和参数量的库，并且能够显示每一层的详细统计信息。

安装 ptflops

你可以使用以下命令安装 ptflops：

pip install ptflops

代码示例

import torch
from ptflops import get_model_complexity_info

# 定义模型
model = MyModel()

# 计算 FLOPs 和参数量
with torch.cuda.device(0):
    flops, params = get_model_complexity_info(model, (3, 32, 32), as_strings=True, print_per_layer_stat=True)

print(f"FLOPs: {flops}, Params: {params}")

输出结果示例

FLOPs: 402.34 MMac, Params: 131.15 M

Code

AI_With_NumPy
此项目汇集了更多AI相关的代码实现，供大家学习使用，欢迎点赞收藏👏🏻

备注

个人水平有限，有问题随时交流~