【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相关的代码实现，供大家学习使用，欢迎点赞收藏👏🏻
备注
个人水平有限，有问题随时交流~
原文地址：https://blog.csdn.net/yiqiedouhao11/article/details/142978359
免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！