3. 非结构化剪枝与稀疏推理
任务说明
本案例针对一个已训练好的两层全连接网络（MNIST手写数字分类），实施非结构化剪枝：将权重中值小于阈值的单个参数置零，生成稀疏权重矩阵；随后利用稀疏矩阵存储与计算（scipy.sparse）完成前向推理，对比稠密与稀疏两种推理路径的输出一致性，并统计稀疏度与推理耗时。
环境依赖
- Python 3.8+
- NumPy
- SciPy
- （可选）scikit-learn 用于加载 MNIST 子集，若不可用可自行准备输入向量
完整实现
import numpy as np
import time
from scipy import sparse
# ---------- 1. 固定随机种子，构造可复现的两层网络 ----------
np.random.seed(0)
input_dim, hidden_dim, output_dim = 784, 128, 10
W1 = np.random.randn(input_dim, hidden_dim).astype(np.float32) * 0.1
b1 = np.zeros(hidden_dim, dtype=np.float32)
W2 = np.random.randn(hidden_dim, output_dim).astype(np.float32) * 0.1
b2 = np.zeros(output_dim, dtype=np.float32)
# ---------- 2. 非结构化剪枝：值阈值置零 ----------
threshold = 0.05
W1_pruned = np.where(np.abs(W1) < threshold, 0.0, W1)
W2_pruned = np.where(np.abs(W2) < threshold, 0.0, W2)
def compute_sparsity(W):
return 1.0 - np.count_nonzero(W) / W.size
sparsity_W1 = compute_sparsity(W1_pruned)
sparsity_W2 = compute_sparsity(W2_pruned)
print(f"W1 稀疏度: {sparsity_W1:.2%}")
print(f"W2 稀疏度: {sparsity_W2:.2%}")
# ---------- 3. 构造测试输入（模拟单张MNIST图像） ----------
x = np.random.randn(input_dim).astype(np.float32)
# ---------- 4. 稠密前向推理（基准） ----------
def dense_forward(x, W1, b1, W2, b2):
h = np.maximum(0, x @ W1 + b1) # ReLU
out = h @ W2 + b2
return out
out_dense = dense_forward(x, W1_pruned, b1, W2_pruned, b2)
# ---------- 5. 稀疏前向推理（CSR格式） ----------
W1_sparse = sparse.csr_matrix(W1_pruned)
W2_sparse = sparse.csr_matrix(W2_pruned)
def sparse_forward(x, W1_sparse, b1, W2_sparse, b2):
h = np.maximum(0, W1_sparse.dot(x) + b1) # 注意：CSR.dot(向量)
out = W2_sparse.dot(h) + b2
return out
out_sparse = sparse_forward(x, W1_sparse, b1, W2_sparse, b2)
# ---------- 6. 验证一致性 ----------
max_diff = np.max(np.abs(out_dense - out_sparse))
print(f"稠密与稀疏输出最大误差: {max_diff:.2e}")
assert max_diff < 1e-4, "稀疏推理结果与稠密不一致！"
# ---------- 7. 推理耗时对比（多次取平均） ----------
def timeit(func, repeat=100):
t0 = time.perf_counter()
for _ in range(repeat):
func()
return (time.perf_counter() - t0) / repeat
t_dense = timeit(lambda: dense_forward(x, W1_pruned, b1, W2_pruned, b2))
t_sparse = timeit(lambda: sparse_forward(x, W1_sparse, b1, W2_sparse, b2))
print(f"稠密推理平均耗时: {t_dense*1e6:.2f} µs")
print(f"稀疏推理平均耗时: {t_sparse*1e6:.2f} µs")
print(f"加速比: {t_dense/t_sparse:.2f}x")运行方式
python prune_sparse_inference.py验证方法与判定标准
- 正确性：程序输出中“稠密与稀疏输出最大误差”应小于
1e-4，且断言通过，证明稀疏路径与稠密路径数学等价。 - 稀疏度：打印的
W1与W2稀疏度应明显大于 0（阈值 0.05 下通常 W1 约 50% 以上、W2 约 40% 以上），说明非结构化剪枝确实产生了稀疏模型。 - 加速效果：观察“加速比”数值。由于本例网络较小且单次推理开销低，加速比可能接近 1 或略小于 1（稀疏矩阵格式转换与函数调用开销占主导）；若需观察明显加速，可增大
hidden_dim至 512/1024 或增加repeat次数。判定通过的标准是：稀疏推理输出与稠密一致，且稀疏度大于 0，而非必须加速比大于 1。
关键说明
- 非结构化剪枝将单个权重置零，不改变矩阵维度，因此稀疏度越高，理论上可利用稀疏矩阵库跳过零值计算。
- 本案例使用
scipy.sparse.csr_matrix的dot方法执行稀疏矩阵-向量乘，这是稀疏推理的核心加速点。 - 实际部署中，稀疏模型还需配合专用硬件或推理库（如 DeepSparse、TensorRT 稀疏模式）才能获得理想加速；本案例聚焦算法正确性与稀疏格式转换路径。