12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091 |
"""演示27: 模型复杂度 — Params 与 FLOPs 的手工计算与对比。
用 numpy 手算常见层的参数量(Params)与浮点运算量(FLOPs):
- 全连接 FC: W∈R^{m×n} -> params=mn+m, flops=2mn·B
- 卷积 Conv: K∈R^{k×k×C_in×C_out} -> params=k²·C_in·C_out(+bias)
- 多头注意力 MHA: 4 个 W∈R^{d×d} -> params=4d²(+位置偏置)
并对比"显存占用(权重字节)"与"不同网络参数量级"。
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from figure_style import figures_dir, configure
configure()
OUT = figures_dir()
os.makedirs(OUT, exist_ok=True)
def fc_params(n_in, n_out):
return n_in * n_out + n_out # W + bias
def fc_flops(n_in, n_out, batch):
return 2 * n_in * n_out * batch # 乘+加
def conv_params(k, c_in, c_out):
return k * k * c_in * c_out + c_out # 核 + bias
def conv_flops(k, c_in, c_out, h, w, batch):
return 2 * k * k * c_in * c_out * h * w * batch # 每个输出像素一次乘加
def mha_params(d_model, num_heads):
return 4 * d_model * d_model # Q/K/V/out 四组 W
# ---- 1. 逐层计算一个小型 CNN 的 Params/FLOPs ----
print("=" * 60)
print("小 CNN 复杂度逐层统计 (输入 3x224x224, batch=1)")
layers = [
("Conv3-64", conv_params(3, 3, 64), conv_flops(3, 3, 64, 224, 224, 1)),
("Conv3-128", conv_params(3, 64, 128), conv_flops(3, 64, 128, 112, 112, 1)),
("Conv3-256", conv_params(3, 128, 256), conv_flops(3, 128, 256, 56, 56, 1)),
("FC-1024", fc_params(256 * 28 * 28, 1024), fc_flops(256 * 28 * 28, 1024, 1)),
("FC-1000", fc_params(1024, 1000), fc_flops(1024, 1000, 1)),
]
total_p = total_f = 0
for name, p, f in layers:
total_p += p; total_f += f
print(f" {name:10s} Params={p:>12,} FLOPs={f:>15,.0f}")
print(f" 合计 Params={total_p:,} ({total_p/1e6:.1f}M), FLOPs={total_f:,.0f} ({total_f/1e9:.1f}G)")
# ---- 2. 显存占用(权重字节, float32=4B) ----
fig, axes = plt.subplots(1, 2, figsize=(12, 4.8))
ax = axes[0]
names = ["FC(1000x1000)", "Conv(3x3, 512→512)", "ResNet-50", "ViT-Base", "GPT-2 1.5B", "GPT-3 175B"]
params = [1e6, 2.36e6, 25e6, 86e6, 1.5e9, 175e9]
bytes_w = [p * 4 for p in params] # float32 权重字节
ax.barh(names, [b / 1e9 for b in bytes_w], color="tab:blue")
ax.set_xlabel("权重显存 (GB, float32)")
ax.set_title("参数量 → 模型文件/显存占用 (4 字节/参数)")
ax.grid(axis="x", alpha=0.3)
# ---- 3. Params 与 FLOPs 量级对比(不同模型族) ----
ax = axes[1]
models = {
"LeNet-5 (1998)": (60e3, 3e8),
"VGG-16 (2014)": (138e6, 1.55e10),
"ResNet-50 (2015)": (25e6, 4.1e9),
"ViT-Base (2020)": (86e6, 1.8e10),
"GPT-3 (2020)": (175e9, 3.1e14),
}
for name, (p, f) in models.items():
ax.scatter(p, f, s=60, label=f"{name}\n({p/1e6:.0f}M, {f/1e12:.2f}T)")
ax.set_xscale("log"); ax.set_yscale("log")
ax.set_xlabel("参数量 Params"); ax.set_ylabel("单样本 FLOPs (log)")
ax.set_title("模型复杂度: Params vs FLOPs (越大越强也越贵)")
ax.legend(fontsize=7)
for axx in axes:
axx.grid(alpha=0.3)
fig.suptitle("模型复杂度评估: 参数量 Params × 计算量 FLOPs", fontsize=14)
fig.tight_layout(rect=(0, 0, 1, 0.93))
out = os.path.join(OUT, "fig_train_params.png")
fig.savefig(out, dpi=150)
print(f"-> {out}")
|
