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

最新推荐文章于 2025-08-09 21:54:05 发布

原创

最新推荐文章于 2025-08-09 21:54:05 发布 · 1.3k 阅读

10 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

在深度学习模型中，评估模型的复杂性和效率时，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(