前言

一般情况下我们都是使用折线图绘制和监控我们的损失函数， y 轴是损失函数的值，x 轴是训练的轮次。 这种情况下我们只有损失函数空间的一维视图，并且只能看到小范围的参数梯度。

有没有一种方法能够让我们将的GPT的1750亿参数损失空间进行可视化呢？如果可以的话，那么这数十亿个参数的梯度会怎样?虽然理论上是可行的，但是这么大的参数对于我们现在的硬件条件来说还是不可行的，但是如果我们至少可以在一个小维度（比如，3维）空间中看到损失图像，这样是否可行呢？这篇介绍性文章简要说明了它是如何实现的，以及它是一个多么简单而又引人入胜的想法。

You can’t step past [someone] in this [two] dimension. Observe this two-dimensional egg. If we were in the third dimension, looking down, we’d be able to see an unhatched chick in it, just as a chick inside a three-dimensional egg could be seen by an observer in the fourth dimension

Prof. Farnsworth, Futurama E15S7

在训练神经网络时，我们绘制的损失函数会根据模型架构、优化器、初始化方法等不同配置而不同。 虽然这些选择对最终目标的影响尚不清楚，但是我们可以将损失函数的收敛进行可视化，这不仅是为了好玩，也是为了深入了解训练的过程以及结果。损失函数三维图的有助于解释为什么神经网络可以优化极其复杂的非凸函数，以及为什么优化的最小值能够很好地被推广。 （例如，通过这种可视化观察到了对残差连接的一个有用的现象：它们可以防止模型将损失三维图变得混乱，因此在训练中很有用 [3]）。

数据集

为简单起见本文使用 MNIST 数据集，并探讨了可视化的几个方面。 更多关于其他数据集和这个想法方面的学术研究可以在关于这个问题的论文中找到 [2, 3, 4],首先让我们配置dataset和dataloader

from torchvision import transforms as T

transform = T.Compose([
    T.ToTensor(),
    T.Resize((28,28)),
    T.Normalize((0.1307,), (0.3081,))
  ])
dataset = Dataset(root='data', download=True, train=False, transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,
                                batch_size=len(dataset),
                                pin_memory=True, num_workers=4)

目标

然后，我们设