BP神经网络3D可视化：从数学原理到Python实现

1次阅读

共计 2586 个字符，预计需要花费 7 分钟才能阅读完成。

训练神经网络就像在黑暗中进行手术——我们输入数据、计算损失、反向传播，但权重矩阵的变化过程始终是个黑盒子。传统训练曲线只能反映整体误差变化，而关键的层间权重分布和梯度流动却难以观察。当遇到模型不收敛时，我们常陷入盲目调参的困境：

是学习率设置不当导致梯度爆炸？
还是权重初始化不合理造成神经元死亡？
或是网络结构设计缺陷引发梯度消失？

虽然 TensorBoard 和 Weights & Biases 提供了出色的训练监控，但在 3D 可视化方面存在明显不足：

展示维度受限：主要支持 2D 投影或切片视图
实时性不足：通常需要完整训练周期后才能查看
交互能力弱：难以动态旋转 / 缩放观察权重分布

import torch
import torch.nn as nn
import numpy as np

# 构建带 3 个隐藏层的测试网络
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 50)  # 故意设计不对称结构
        self.fc2 = nn.Linear(50, 30)
        self.fc3 = nn.Linear(30, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))
        return self.fc3(x)

# 生成螺旋线测试数据
def generate_data(n=1000):
    theta = np.linspace(0, 4*np.pi, n)
    r = np.linspace(0, 5, n)
    x = r * np.sin(theta)
    y = r * np.cos(theta)
    z = np.sin(2*theta) + np.random.normal(0, 0.1, n)
    return torch.FloatTensor(np.column_stack([x,y])), torch.FloatTensor(z)

# 注册前向 / 反向 hook 存储权重和梯度
weight_maps = {}
gradient_maps = {}

def register_hooks(model):
    def hook_fn(module, input, output, name):
        if isinstance(module, nn.Linear):
            weight_maps[name] = module.weight.detach().numpy()

    def back_hook_fn(module, grad_input, grad_output, name):
        if isinstance(module, nn.Linear):
            gradient_maps[name] = module.weight.grad.detach().numpy()

    for name, module in model.named_modules():
        module.register_forward_hook(lambda m, i, o, n=name: hook_fn(m, i, o, n))
        module.register_full_backward_hook(lambda m, gi, go, n=name: back_hook_fn(m, gi, go, n))

from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

def update_plot(epoch):
    ax.clear()

    # 绘制权重分布
    for i, (name, weights) in enumerate(weight_maps.items()):
        xs = weights.flatten()
        ys = np.full_like(xs, i)  # 层编号作为 Y 轴
        zs = np.random.uniform(-1,1,len(xs))  # 添加随机深度
        ax.scatter(xs, ys, zs, label=f'{name} weights')

    # 绘制梯度箭头
    if gradient_maps:
        for i, (name, grads) in enumerate(gradient_maps.items()):
            arrows = grads.flatten()[:50]  # 抽样显示
            starts = weight_maps[name].flatten()[:50]
            ax.quiver(starts, np.full_like(starts, i), 
                np.zeros_like(starts), -arrows, 
                np.zeros_like(starts), length=0.5, 
                color='r', alpha=0.3)

    ax.set_xlabel('Weight Value')
    ax.set_ylabel('Layer Depth')
    ax.set_zlabel('Random Spread')
    ax.set_title(f'Epoch {epoch}')
    plt.legend()
    plt.pause(0.1)  # 控制刷新率