📚 课程目录第2章 ONNX模型导出与结构可视化

2. ONNX模型导出与结构可视化

任务目标

分析训练好的PyTorch模型各卷积层权重的数值分布，绘制直方图并计算关键统计量（均值、标准差、最大值），为后续量化时选择缩放因子和零点提供数据依据。

环境依赖

Python 3.8+，PyTorch 1.10+，Matplotlib 3.5+，NumPy。无需额外数据集，使用随机初始化的模型即可完成分布分析流程。

核心实现

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

# 定义一个包含不同层类型的小型CNN，用于演示权重分布分析
class DemoCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc = nn.Linear(32 * 8 * 8, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        return self.fc(x)

def analyze_weight_distribution(model):
    """遍历模型所有带权重的层，返回每层权重张量及其名称"""
    weight_layers = []
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:  # 排除bias和1维参数
            weight_layers.append((name, param.detach().cpu().numpy()))
    return weight_layers

def plot_weight_histograms(weight_layers, bins=50, figsize=(12, 4 * len(weight_layers))):
    """为每层绘制权重直方图，并标注统计量"""
    n_layers = len(weight_layers)
    fig, axes = plt.subplots(n_layers, 1, figsize=figsize)
    if n_layers == 1:
        axes = [axes]

    for ax, (name, weights) in zip(axes, weight_layers):
        # 展平权重值用于直方图
        flat = weights.flatten()
        ax.hist(flat, bins=bins, density=True, alpha=0.7, color='steelblue', edgecolor='black')
        ax.set_title(f'{name}\nshape={weights.shape}', fontsize=10)

        # 标注统计信息
        mean = np.mean(flat)
        std = np.std(flat)
        max_abs = np.max(np.abs(flat))
        stats_text = f'mean={mean:.4f}\nstd={std:.4f}\nmax|w|={max_abs:.4f}'
        ax.text(0.98, 0.95, stats_text, transform=ax.transAxes,
                ha='right', va='top', fontsize=9,
                bbox=dict(boxstyle='round', facecolor='lightyellow', alpha=0.8))
        ax.set_xlabel('Weight value')
        ax.set_ylabel('Density')

    plt.tight_layout()
    return fig

# 主流程
if __name__ == '__main__':
    torch.manual_seed(42)  # 固定随机种子保证可复现
    model = DemoCNN()

    # 可选：加载真实训练权重，此处用随机初始化演示
    # model.load_state_dict(torch.load('your_model.pth'))

    layers = analyze_weight_distribution(model)
    print(f'发现 {len(layers)} 个权重层')
    for name, w in layers:
        print(f'{name}: shape={w.shape}, range=[{w.min():.4f}, {w.max():.4f}]')

    fig = plot_weight_histograms(layers)
    plt.savefig('weight_distribution.png', dpi=150, bbox_inches='tight')
    plt.show()

运行方式

将上述代码保存为 weight_hist.py；
在终端执行 python weight_hist.py；
程序会打印每层权重张量的形状与数值范围，并弹出直方图窗口，同时保存为 weight_distribution.png。

验证方法与判定标准

输出检查：控制台应显示3个权重层（conv1.weight、conv2.weight、fc.weight），每层数值范围大致在[-0.5, 0.5]内（PyTorch默认初始化）；
图形检查：每个子图应呈现以0为中心的单峰分布，且conv层（3×3卷积核）的分布比fc层更集中——因为卷积核共享权重，参数量少但分布更紧凑；
量化依据：观察各层max|w|值，该值即为对称量化时选择缩放因子的直接参考；若后续采用非对称量化，则需结合min和max确定零点。

量化参数选择指引

直方图显示权重近似对称分布时，优先采用对称量化（零点=0），缩放因子 s = max|w| / 127（int8）。若分布明显偏移（均值远离0），则需使用非对称量化，此时 s = (max - min) / 255，零点 z = -round(min / s)。本案例中随机初始化的权重均值为0，适合对称量化；实际训练后的模型需按此流程重新分析。