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

1次阅读

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

在训练神经网络时，我们常常遇到这样的困扰：模型突然停止收敛、验证集准确率波动异常，却难以快速定位问题根源。传统的训练日志和 2D 曲线只能反映整体 loss 变化，而真正影响模型表现的隐藏层权重分布、梯度流动等关键信息，却像黑盒子一样难以直观观察。

TensorBoard：擅长展示标量指标和计算图，但 3D 可视化需要额外插件
Weights & Biases：云端的实验跟踪利器，但对实时 3D 渲染支持有限
Netron：网络结构查看工具，不涉及训练过程动态

这些工具的共同短板是：无法在训练过程中实时观察高维参数的 空间分布变化。

import torch
import torch.nn as nn

class ThreeLayerNet(nn.Module):
    def __init__(self, input_size=10, hidden_size=50):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.out = nn.Linear(hidden_size, 1)

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

def register_hooks(model):
    activations = {}
    gradients = {}

    def hook_fn_forward(module, input, output, name):
        activations[name] = output.detach()

    def hook_fn_backward(module, grad_input, grad_output, name):
        gradients[name] = grad_output[0].detach()

    # 为每个全连接层注册钩子
    model.fc1.register_forward_hook(lambda m, i, o: hook_fn_forward(m, i, o, 'fc1'))
    model.fc2.register_forward_hook(lambda m, i, o: hook_fn_forward(m, i, o, 'fc2'))

    model.fc1.register_full_backward_hook(lambda m, gi, go: hook_fn_backward(m, gi, go, 'fc1'))
    model.fc2.register_full_backward_hook(lambda m, gi, go: hook_fn_backward(m, gi, go, 'fc2'))

    return activations, gradients

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

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

def update_plot(epoch):
    ax.clear()

    # 绘制权重分布
    weights = model.fc1.weight.data.numpy()
    ax.scatter(weights[:,0], weights[:,1], weights[:,2], 
               c='r', marker='o', label='Weights')

    # 绘制梯度方向
    if 'fc1' in gradients:
        grads = gradients['fc1'].numpy()
        ax.quiver(weights[:,0], weights[:,1], weights[:,2],
                 grads[:,0], grads[:,1], grads[:,2], 
                 length=0.1, color='b', label='Gradients')

    ax.set_xlabel('X Axis')
    ax.set_ylabel('Y Axis')
    ax.set_zlabel('Z Axis')
    ax.view_init(elev=20, azim=epoch*2)  # 动态旋转视角

ani = FuncAnimation(fig, update_plot, frames=100, interval=200)
plt.show()