1. 基于PyTorch的模型参数统计与FLOPs计算

任务目标

为后续模型裁剪建立量化基准：编写一个可复用的PyTorch脚本，统计给定CNN模型的总参数量、可训练参数量以及理论FLOPs（浮点运算次数）。本案例以轻量级CNN（类似VGG风格的简化网络）为例，输出结构化的统计报告。

环境与依赖

  • Python 3.8+
  • PyTorch 1.12+（含torchvision，用于加载预训练模型或自定义模型）
  • thop库（用于FLOPs计算）：pip install thop

关键完整代码

import torch
import torch.nn as nn
from thop import profile, clever_format

# ---------- 1. 定义示例CNN模型（可替换为任意nn.Module） ----------
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),   # 输出: 32x32x32
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),                              # 输出: 32x16x16
            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # 输出: 64x16x16
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),                              # 输出: 64x8x8
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 8 * 8, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

# ---------- 2. 参数统计函数 ----------
def count_parameters(model):
    """返回总参数量与可训练参数量（单位：百万）"""
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total / 1e6, trainable / 1e6

# ---------- 3. FLOPs统计函数 ----------
def compute_flops(model, input_size=(1, 3, 32, 32)):
    """使用thop计算FLOPs与参数量（thop内部也统计参数）"""
    device = next(model.parameters()).device
    dummy_input = torch.randn(input_size).to(device)
    flops, params = profile(model, inputs=(dummy_input,), verbose=False)
    return flops, params

# ---------- 4. 主执行入口 ----------
if __name__ == "__main__":
    # 初始化模型
    model = SimpleCNN(num_classes=10)
    model.eval()

    # 参数统计
    total_m, trainable_m = count_parameters(model)
    print(f"总参数量: {total_m:.4f} M")
    print(f"可训练参数量: {trainable_m:.4f} M")

    # FLOPs统计（输入尺寸需与模型实际输入匹配）
    flops, params = compute_flops(model, input_size=(1, 3, 32, 32))
    flops_readable, params_readable = clever_format([flops, params], "%.3f")
    print(f"FLOPs: {flops_readable}")
    print(f"thop统计参数量: {params_readable}")

    # 额外：打印每层参数量（便于定位裁剪目标）
    print("\n各层参数量明细：")
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(f"{name}: {param.numel():,} 参数")

运行与验证

  1. 运行方式：将上述代码保存为model_stats.py，在终端执行python model_stats.py。
  2. 预期输出（基于代码逻辑直接推算，非实测值）：
    • 总参数量与可训练参数量相等（本模型无冻结层），约为 0.21 M（具体数值由网络结构决定）。
    • FLOPs约为 0.09 G（输入3×32×32时）。
    • 每层参数明细会列出features.0.weight、classifier.1.weight等各层参数数量。
  3. 通过条件：
    • 脚本无报错运行完成；
    • 输出的参数量与FLOPs为有限正数；
    • 若替换为其他CNN（如torchvision.models.resnet18），只需修改模型初始化部分并调整输入尺寸，脚本仍能正确输出统计结果。

关键说明

  • FLOPs计算原理：thop通过向模型传入一个随机输入，跟踪各层实际执行的乘加运算次数。因此输入尺寸必须与模型期望尺寸一致，否则会报错或得到错误结果。
  • 参数量与FLOPs的区别：参数量反映模型存储开销，FLOPs反映单次推理的计算量。裁剪时通常同时关注两者：参数量影响模型文件大小，FLOPs影响推理延迟。
  • 扩展性：本脚本中的SimpleCNN可替换为任何继承自nn.Module的网络（包括加载的预训练模型），统计函数无需修改。
← 上一章 📖 返回目录 下一章 →