前言:
🧠 本系列的代码工具和图表都是作者在实战中总结下来的,适合需要从训练监控中挖掘问题线索的开发者。想系统学习这类监控方案的朋友,欢迎评论区回复【训练监控】。
正文:
这是一系列文章的第一篇,教你如何提升神经网络训练过程中的监控和排查能力:
- 更好的方式来监控神经网络训练(本篇)
- 神经元死亡问题(即将推出)
- 梯度消失与爆炸(即将推出)
- 小心梯度震荡(即将推出)
所有图片均由作者提供
有时候让你的人工神经网络表现得好并不容易。训练中卡住了,或者出问题了,找出原因真的挺难的。这时候你需要好工具。能把底层问题勾勒出来的工具。能在训练时可视化模型内部运作的工具。是那种能帮你迅速找出问题根源的工具。
我学习过程中最重要的经验之一,就是从开发早期就开始暴露模型内部的各种指标。信息暴露得越早,开发周期越短,问题发现得越快。
但我感觉我们整个圈子里,现在都还只是在用最基础的监控工具。很多时候,大家就是盯着 loss 和一个跟任务相关的准确率版本,拿这些东西来评估模型结构和训练方案的效果。
像 TensorBoard 和 Weights & Biases 这样的工具确实有帮助,但它们还可以做得更好。这些工具主打的就是能快速生成常规的 loss 和准确率图,并查看模型输出。它们也加入了权重的直方图,这是好事,但除此之外分析能力并不多。它们主要追求的是易用性。这当然重要,但代价就是只能提供比较浅层的信息。虽然它们支持你自己加插件或数据收集机制,但大多数人可能就只用默认的功能——因为你还能想画点啥呢?
这篇文章就是我准备要回答这个问题的开始——还有什么别的东西值得画?我也会介绍一堆有用的指标和值,能帮助你更深入地理解模型动态,以及怎么收集这些数据。我希望你能因此增加几把工具,让你更早发现问题、更快修模型,而且能掌控得比你想象得还深。
比如说这个仪表盘,就能让你在训练时看到模型内部行为的方方面面:
💡 这类仪表盘图是作者经常用的调试方法,在芯片+模型协同优化任务中非常关键,很多隐藏问题靠它能提前揪出来。
训练时可见模型动态细节,能更早发现问题
这只是一个预告,后面还有更多背景知识。等你读完,希望你也能动手搞个自己的版本。
目录
- 问题在哪里
- 解决思路
- 更好的 loss 和 metric 图
- 权重可视化
- 所有层的分布图
- 各层输出的分布图
- 神经元激活率和死亡率可视化
- 尺度对比
- 模型内省
- 数据采集方法
- 一页全览
- 总结
- 源代码
📦 本文所有图表都支持复现,已整理成 notebook + 可复用 toolkit,评论区或者私了解获取项目的方法,这些方法可以直接套用到你的项目中。
问题在哪里
在讲怎么解决之前,我先详细讲讲问题到底出在哪。
神经网络训练过程里可能遇到各种问题。从代码小错误,到经典的梯度消失、学习率过大,再到更复杂更隐蔽的问题,比如神经元死亡。这些在学术界早就讨论很多了,也有不少办法来测量和可视化模型的行为和这些问题的影响。常见技术包括:计算权重或梯度的均值和标准差、欧几里得范数、量之间的时间相关性、PCA降维成2D可视化、数据流聚类等等。
这些技术很强,但不太好用。很多人也压根不知道它们,或者不知道它们对实际神经网络设计和排障有多大帮助。
比如说你在用 TensorFlow,第一个问题就是它压根不提供访问梯度的接口。你甚至不能用自定义 Callback 来搞。你得自己写个自定义训练循环。而在 PyTorch 里,这种自定义训练循环倒是比较常见,这步稍微轻松点。
第二个问题是:怎么收集你需要的数据,又不让训练变慢或爆内存?而且还得能把这些都可视化出来。通常权重数组是多维的张量,值成千上万,一个模型里有很多这样的权重。
解决思路
在这篇文章里,我会介绍几种可视化方法,能在训练过程中让你更清楚地看到模型内部发生了什么。我也会讲一些怎么收集数据的思路。
我用的是 matplotlib 来画图,够简单,不像 TensorBoard 那样还有渲染限制,结果也更好复现。你可以试试怎么把这些图集成进你喜欢的平台里。
为了方便你快速用起来,我把模型内省和可视化的工具打包做成了一个 toolkit,并开源在 GitHub 上了。想快速上手的可以直接插进训练流程里用。这套东西是用 Callback 风格写的,TensorFlow 用户应该很熟悉。
本文中所有的可视化都是用这个 toolkit 做的。但我要分享的点子和技术跟 toolkit 本身没啥强依赖,所以这里不会赘述它的细节。
🔧 这套 toolkit 是作者在部署推理引擎+调试神经元死亡问题时做的内部工具,后来提炼开放出来了,欢迎尝试集成到你的训练流程里。
那我们就开始吧。
更好的 Loss 和 Metric 图
我们在开发神经网络模型时,常常会先训练几轮短期跑一下看看有没有基本起效。问题是,很多细微的问题不会那么快暴露出来,只有长训练后才会出事。如果你只看 loss 曲线,这些问题一开始根本看不出来。但我们可以换个方式。
我用两个相似模型跑了个简单的分类问题。下面是它们前几轮训练的 loss 曲线:
相似的模型,前几轮训练 loss 曲线也很接近
你看起来会觉得两个模型训练得都还行,几乎没啥差别。甚至我还可能猜 Model B 比 A 更好点。现在只训练了6轮确实有点假设性强,但这类教训在实际场景里也适用。我们来看它们训练久一点后的结果:
长周期训练后的 loss 曲线
现在差别一下子就大了。Model B 出现了梯度震荡的迹象。其实也不奇怪,因为真相是:这两个模型其实是一样的,唯一的区别是 Model A 的学习率合适,而 Model B 的学习率稍微大了一点点——不至于一开始就出问题,但靠近最优点时就崩了。
TensorFlow 的训练历史记录是现成的,loss 和其它指标都会记录下来,还挺好用,可以很容易画出上面的图。但默认的记录粒度是按 epoch 的。每轮训练结束采样一次,通常是该轮 loss 的平均值。
这种做法确实简洁——每轮里 mini-batch 的数据不同,要是每步都记录,那会非常随机,还可能引入噪声。如果有验证集,也没必要每步都跑验证。但问题就是我们错过了训练过程中大量有用信息。
下面这几张图是用相同训练跑出来的,不过这次用了每步的 loss 值。第一行图仍然按 epoch 汇总,但加了百分位数分布:中位数(实线),±25百分位,还有最小最大值区间。第二行图是每步的原始数据。
用每步数据生成的 loss 曲线
现在两个学习率的影响就一目了然了。Model A 图里那些小锯齿,就是 epoch 边界的记录现象,这是 TensorFlow 计算平均值的副作用。所以最终按 epoch 看的时候分布非常窄(方差小)。Model B 的曲线就乱得多,反映在 epoch 曲线里也是分布很宽、很吵。
就算只看上面那一行,我们也能比传统 loss 曲线更早发现潜在问题。我们再来看看文章开头那两个模型的前6轮图,加上百分位分布后再比较一下:
这种方法对其它输出指标也同样有效,比如准确率。
在 TensorFlow 里,要收集每步的 loss 和指标,只需要写一个自定义版本的 History Callback。toolkit 里也有例子可以直接参考。
可视化权重
根据你用的初始化方案不同,模型里的权重通常会初始化成像下面这些分布:
示例:一个 100x100 权重矩阵初始化后的分布图。具体的尺度会随着矩阵形状和尺寸变化而不同。
有些层在训练过程中会大致保留原来的分布,有些则会发生很大偏移。一个状态良好的模型里,权重应该大致分布在 0 的两边,保持平衡。但有时候并不是这样,这时候能获取这类信息就很有用。而且,如果能比较不同层之间、以及不同时间点上的权重尺度,也很实用。
一种很常见的可视化权重和偏置的方式,是把多个直方图叠在一起,形成一个 3D 的层叠图。越旧的状态排在后面,越新的排在前面:
TensorBoard 的 3D 直方图(由 TensorBoard 生成)
看起来确实很漂亮,而且这基本上是目前唯一能表现时间动态直方图的方法。但它并不适合所有场景。
一旦你想用直方图来可视化梯度,就容易出问题。下面是来自两个不同模型的两个代表性层的权重和梯度分布图。最左边是初始化后(用 He-normal 初始化)的权重分布,中间是训练结束时的权重分布,最右边是训练结束时的梯度分布:
你会发现,权重分布可以保留原始的正态曲线,甚至还会更理想;但梯度就不一样了,它们可能非常尖锐、跨度很大。它们根本不遵循正态分布,而且差异极端到直方图里根本看不出啥有意义的东西。简单说:这种情况必须引入对数坐标轴才行。
下面是三种替代方式,用来表示权重、梯度或其他张量的分布:
三种可视化大型张量的方式
这些图来自同一个模型的同一层,展示了训练过程中分布是怎么变化的。第一行是权重,第二行是梯度。
这三种图的含义如下:
- 数值分布(也叫原始分布) :展示的是原始值的各个百分位 —— 中位数是中间那根实线,灰色代表最小/最大区间,其他色带表示中间的其他分位。这种图能很好地显示整体的特征。但它不擅长识别是否有多个波峰 —— 这种情况还是得用详细的直方图来看。另外,比如说梯度的数值分布虽然能显示最小最大范围,但其他信息就没啥看头了。
- 对数量级分布:中间那列图用的是数值的绝对值,再转成对数坐标轴来做百分位。这种方式特别适合看梯度,有时候也适用于权重或别的张量。因为对数取的是绝对值,符号信息会丢失,而你有时是关心正负是否平衡的,所以用橙色线来标示“质量分布中心”在哪里。0% 就表示正负完全平衡,+100% 就是全是正值,-100% 就是全是负值。
- 范数图:右边那列图用了另一个在学术圈很常见的指标 —— L2 范数,也叫欧几里得范数。它是向量范数在多维空间的扩展 —— 计算方式是所有元素平方和的平方根。你可以把整个张量当成一个多维向量,然后算它的“长度”。图里的范数值已经调整成跟张量尺寸无关了。这样我们就可以看出这个张量里的数值整体上有多大。从图中可以看出,权重的尺度略微上升,而梯度的尺度则降到了原来的 40%。
每张图中间的那行字就是对应张量的 shape,方便你解读。小张量(比如偏置向量、最后全连接层)往往会表现出更混乱的分布,用这个信息能帮上忙。
在这三种方式里,我发现最实用的是前两个分布图。到底选哪个得看你数据的具体情况,有时候也得多试几种。比如你上面看到的,原始数值分布对权重来说效果就很好。
但对于梯度来说,对数量级分布是必须的。因为它能更好地处理梯度那种天生就跨度巨大的情况。你既能看全梯度的变化范围,又能发现梯度是不是小到几乎没用了。上面那个梯度图显示了 9 个百分位 —— 从 0 到 100% 均匀划分,每个间隔 12.5%。注意那个深蓝色阴影,在大约第 18 个 epoch 时掉到底部,说明有 37.5% 的梯度已经变成 0 了。而到了第 47 个 epoch,有 50% 的梯度为 0。这可能是在提示你出了问题。
范数图适合你只想要一个标量值来代表某层的权重、梯度或其他数据时。比如说要比较多个层之间的差别,这个就很实用。但 L2 范数也有个问题,就是它很容易被张量尺寸影响。比如一个全是 1 的张量,它的平方和就等于张量的元素总数。所以我更喜欢稍微改良一下 —— 把 L2 范数除以张量大小的平方根。这个方法其实就等价于算这个张量的 RMS(均方根)。有了这个“尺寸归一化”的范数,你就能比较不同尺寸的层之间的梯度流动情况了。
所有层的分布图
不管你喜欢用哪种方法,接下来把它们画到模型里每个可训练参数张量上是很有用的。比如,下面这两张图展示了整个模型在训练过程中的参数分布(用的是数值分布)和梯度分布(用的是对数量级分布):
模型参数在整个训练过程中的分布
梯度在整个训练过程中的分布
这两张图的最上方是两个大的总览图,提供了所有层的整体概况,重点在尺度上。左上角的图展示了所有层 log-数量级中位数的百分位分布范围(无论下面的图是画的原值还是数量级,这个图都更适合看尺度);右上角的图是把各层的尺度叠加在一起比较。我稍后会解释这个图的用意。
我特别喜欢把所有层集中画在一张图里,这样你就能在一个紧凑视图里看到所有内容。不会有那种“点开这个点开那个”的过程,然后你还可能错过一个关键的 toggle box 里面藏着的重要信息。
有了这两组图(每一层的参数和梯度),你就能掌握模型在训练过程中的大量动态信息。这些信息是光靠 loss 曲线看不到的。
参数分布图尤其适合用来判断你是否需要做权重正则化。很多网络就算没有显式的正则项也能训练得不错。这些图能立刻告诉你某些层的权重是不是太大了。
在权重和梯度之间,梯度的可视化通常是最重要的。虽然很多机器学习入门教材会说“权重太大或太小”,但实际上就算不加正则化,这种情况也不常见。根据我的经验,真正频繁出问题的是梯度。“梯度消失和爆炸”——是梯度。“梯度震荡”——还是梯度。“神经元死亡”——本质上也是通过梯度最容易识别的。
层输出的分布图
有时候,查看每一层的输出也挺有用的。
传统上,做计算机视觉的研究人员会看各隐藏层的特征嵌入,用它们来理解模型是怎么工作的。对于带注意力机制的模型,你可以通过注意力矩阵来看模型关注了哪里。
但如果你没有更合适的领域表示方法去理解隐藏层输出呢?一个选择是把前面讲过的那些分布图方法,用在层输出上。
比如,下面就是前一节用的模型和训练过程中的层输出分布图:
各层输出的分布图
你会注意到 ReLU 激活函数的效果很明显,大多数层只输出正数。不同层之间的输出值尺度也不一样,如果尺度太大,这可能是个问题。同时你还能直观看到 batch normalization 层的影响。虽然它的输入都是正数,但它的输出大致在正负之间平均分布。你甚至可以通过对比有无 batch-norm 层的情况,看出它是怎么影响后面 dense 层的。如果你只是因为别人用了 batch-norm 你也跟着加,结果又不看它实际怎么影响模型,那你就是在“摸黑操作”。
我第一次开始去看层输出,是在研究神经元死亡的时候。训练时采集层输出的数据确实比采集梯度麻烦点,所以很多人可能不会第一时间想到这一步。但我很快发现,层输出其实和梯度一样重要,它们对你理解模型训练过程中的行为非常关键。一个原因是:我们在堆层、选激活函数、插 normalization 时,设计的其实就是这些层输出。所以你当然得看看你实际堆出来的是什么效果。比如你只是在别人成功的模型上盲目抄个 batch-norm,但你不看它有没有起作用,那你就是盲调。
如果让我从权重、梯度、层输出这三者中选两个来可视化,我会选梯度和层输出。
那层输出这么重要,它们的梯度呢?
“啊?但你不是已经讲了梯度了吗?”我听到你这么说了。
别急。反向传播算法对每一层是分两步计算梯度的。我们平常关心的,是第二步,也就是计算权重、偏置和其他参数的梯度。但在那之前还有一步,就是计算“关于层输出的梯度”。
没错,我们也可以用相同的可视化方法,来画层输出的梯度分布图:
层输出的梯度分布图
层输出和它们的梯度,其行为和参数及其梯度很类似。输出通常可以直接用原始数值分布来画,虽然有时候用对数量级图会更合适。输出梯度则是又尖又宽 —— 一定得用对数量级分布来画。下面是这三种图分别用在一个典型层上的对比:
层输出和梯度的绘制类型的对比
那输出梯度有什么用?这个我说不出一个特别明确的答案。一般来说,我猜它们在某些情况下对排错可能有用,但我不确定具体是哪些情况。不过这依然是一个我希望你手里要有的工具。它到底能帮你什么,是你自己去发现的事。
在检查神经元死亡时,输出梯度的统计分析会更有价值。
可视化神经元激活率和死亡率
ReLU 层有个老问题:它们的单元可能会“死掉”。
这就是指有些单元开始输出恒定值(通常是0),或者它们的梯度始终为0。这是 ReLU 激活函数一个很有名的问题,也因此才有人提出了各种带泄漏或平滑的 ReLU 替代方案。但很少有人去深入讨论如何检测和监控神经元死亡的问题。
我会在后面的文章里详细解释神经元死亡的原理、成因、解决方法,以及怎么去测。但现在,我先介绍两个受这个问题启发的指标,同时也展示我最喜欢的一种图。
我常用的两个指标是:
- 死亡率
- 激活率
这两个指标都是从单元级别去测的。比如对 ReLU 单元来说,输出非0时算“激活”,输出0时就是“未激活”。而一个“死亡”的 ReLU 单元是:对所有样本都始终输出0。死亡率指的是某一层中这些“死掉”的单元所占比例。
但一个没死的 ReLU 单元,也可能只对一部分样本有反应,所以某一层 ReLU 的激活率通常不是100%。激活率就是:对所有样本平均下来,有多少比例的单元是活跃的。这是对死亡率的一个很有价值的补充 —— 因为低激活率可能在提示你,有些单元快死了。
把这个概念泛化一下,你就可以用“张量中0值的比例”来衡量任意张量的激活率和死亡率。这不仅能用在层输出上,还能用在模型参数、它们的梯度,甚至是层输出的梯度上。这在所有层类型上都适用 —— 卷积、循环、transformer 都可以 —— 因为这些结构本质上都是把相同的单元反复应用在不同输入上,跟一个 batch 的样本是一个道理。
所以我们有如下通用定义:
- 死亡率 —— 在所有样本、空间位置、时间位置和/或输入通道上,输出全为0的通道所占比例
- 激活率 —— 所有维度上非0值的比例
这个概念也可以推广到其它激活函数,只是需要额外处理。
下面这个就是我说的“活跃图”。它显示了每一层的激活率和死亡率,还有模型整体的总览图。蓝色实心表示激活率,红色实心表示单元/通道死亡率。还包括了“空间死亡率” —— 是死亡率的一个变种,我会在后面文章里讲。每个小图里面的标签标出了张量的 shape(这个例子是 BHWC 顺序)以及最终 epoch 时的死亡率。最上方的模型总览图里包含了所有层的平均激活率和死亡率,还有它们的最大最小值范围。
各层输出的活跃图
这张图是我在训练一个卷积分类模型(用的是 CIFAR-10 数据集)时画出来的。
你会注意到,很多层的激活率都低于50%。这是 ReLU 层输出的典型特征。每个单元对一个输入有50%的概率输出非0,所以平均下来大概就在那附近。它们低于50%也可能是在说明模型正在学会区分不同样本。
大多数层的死亡率接近0,但某些层会高一些。有两层在训练早期就达到了50%的死亡率,但后来逐渐恢复了。这可能说明很多事情。一方面可能是模型结构有系统性的问题;另一方面也可能是这些层设计得太大,有冗余,可以剪掉提高效率。
上面这张是针对层输出画的活跃图。当你第一次开始学习神经元死亡时,这种图非常管用。因为神经元死亡最直观的定义就是“神经元输出为0”。你也可以用这张图来直接观察 Dropout 的效果 —— Dropout 层的激活率应略低于前一层,而 Dropout 之后的层的死亡率应该低于不加 Dropout 的情况。
📊 活跃图是 toolkit 中作者最推荐的组件之一,适用于 ReLU 死亡监控、Dropout 分析等情境。图表生成逻辑和图例参数也都开放了,想复用可以看 toolkit 或评论区留言或私信。
不过训练时采集层输出要比采集梯度麻烦得多。而实际证明,画梯度的活跃图对理解神经元死亡对训练的影响更有帮助。下面这张图就是同一个模型同一轮训练,对梯度画的活跃图:
模型参数梯度的活跃图
我觉得这张图更能体现模型潜在问题。因为在一个表现良好的模型里,梯度的激活率应该始终接近100%。如果这里掉得很厉害,或者很多层激活率都很低,那说明梯度在某处受限了,模型学习速度慢。
就像我们前面在讲分布图时发现的那样,活跃图也可以应用在模型参数和层输出的梯度上。对模型参数来说,这方法意义不大 —— 因为参数几乎不会正好变成0。但对输出梯度来说,这提供了关于神经元死亡的另一个角度 —— 从“反向传播的梯度是怎么影响这一层”的角度来看。
上面举的这个模型在任务上表现还是不错的,尽管它的神经元死亡率看起来很高。但这不能当作“典型例子”。你会发现,不同的模型结构、超参数和任务领域,会有非常多种不同的行为。
最大的好处是:你能看到这些指标,就能在尝试不同模型时做横向比较 —— 不同模型大小、正则方式、激活函数,都会直接反映在这些图上。这些图会在你没意识到模型表现不好时给出信号,也能提醒你模型是不是太大、该优化了。
说白了,除了最简单的神经网络,神经元死亡是几乎所有模型都会碰到的问题。无论你用不用 ReLU。你只是一直没看见它而已 —— 因为你之前没有这些可视化工具。
比较尺度
我在这里讲的所有数据收集和可视化技术,都是为了让你更容易发现常见问题。有一类常见问题和模型参数及其梯度的尺度有关。最有名的两个就是“梯度消失”和“梯度爆炸”。梯度消失时,前面的层的梯度比后面的层小;梯度爆炸时,所有层的梯度都非常大,远远超过 1.0。但还有别的尺度问题,比如权重变得过于小或过于大。
前面讲的分布图适合用来看所有层的整体尺度变化,以及它们各自分布的范围。这些图能很好地帮你识别比如“全部层都往错误方向偏移”这种情况,比如梯度爆炸。
但还有一种不同的图,是专门看各层之间的相对尺度。这种更适合发现梯度消失这种问题——因为重点是层与层之间的差异。
下面是一个常见的图,用来检测梯度消失:
层尺度图
底部列出了各个层,在 y 轴上画的是某种尺度指标(这里我用的是“尺寸归一化范数”)。为了避免梯度消失,最关键的是第一层的梯度尺度不能小于最后一层。你在图里也能看到,确实是这样,所以暂时没问题。
不过这个图有个问题:当你加上更多 epoch 后,它会很快变得混乱、不好读。上面这个图只用了训练过程中的 5 个 epoch,就已经给出了一个很好的总览,能让你判断有没有梯度消失的问题。
但我更喜欢那些不只对一个问题有用的可视化方法。如果能对每个 epoch 都表示出这个信息就更好了。
这是我后来开始用的另一种方法:
层尺度堆叠图
乍一看这图可能有点奇怪。它的想法是:别管数学准不准,先抓住直觉。每一条颜色交替的带状条代表一层。条带越粗,说明该层的梯度越大。高度是对数量级的(log scale),所以最细的条带比最粗的差了好几个数量级。你可以清楚地看到,最底层(第一层)的梯度比最顶层(最后一层)大,所以不存在梯度消失。而且你也可以看到这个趋势在整个训练过程里都挺稳定。
这张图的步骤如下:
- 计算每一层的尺度指标(比如:尺寸归一化范数)
- 对这些指标取对数
- 从所有 log 值中减去最小值再减1,这样就能保证所有值都是正的,最小值为1
- 重新缩放,使它们加起来等于1.0
- 每个 epoch 分开算
- 用 matplotlib.stackplot() 来画
这个处理方式相当于把最小值当作单位基准,每增加一个倍数就代表尺度上升一个数量级。示例代码可以在项目里找到。
这张图特别适合看梯度震荡的问题。如果一切正常,你会看到图带很平滑;如果有震荡,图会变得很锯齿:
梯度震荡图
🧠 这张“梯度震荡图”是作者在做过的一个 Transformer 微调项目里救过命的 —— 本来以为是模型崩了,结果是梯度异常。后来几乎所有项目都配上这图。
另外在模型严重崩溃的时候也很好用。比如下面这个例子,是因为神经元死亡太多导致的梯度崩溃:
神经元死亡导致的模型崩溃图
在选择“梯度尺度”指标时,有几种效果相近的选项:尺寸归一化的 L2 范数(也叫 RMS)、平均绝对值、原始数值的标准差,这些都能正常使用,也能处理各种边缘情况。
但原始值的平均值不太适合作为尺度指标,它只是衡量偏离0的程度。你如果已经有了百分位数据,那么 25% 和 75% 之间的距离一半,其实就和“平均绝对值”差不多。但这在遇到神经元大量死亡等边缘情况时容易失效。当然,如果你不专门测神经元死亡,那这种“易崩溃”反而是个优势。
模型内省
我前面介绍了很多种不同的可视化方式,也说了它们在不同常见问题中的用法。但这还只是开始。我要讲的是一种更广义的能力 —— 通过修改你的训练过程来收集更多信息,从而针对你具体的问题或领域去定制可视化手段。
要做到这一点,你得有一种高效又通用的方式来收集数据,并且能复用到不同的模型上。
🚀 神经网络训练失败往往不是某一个点的问题,而是整套系统设计的反馈不够可见。这也是作者在项目中必须做训练内省的主要原因。
除了最基础的 loss、accuracy 等指标,还有四组和模型相关的数据在本文中被提到。我把它们称为四种“视角”,因为它们从不同角度告诉你模型内部的情况:
- 可训练参数(也叫变量)—— 权重、偏置等
- 梯度(针对上面的可训练参数)
- 层输出
- 层输出的梯度
但这些信息不是随手就能拿到的。如果你是在跟着常规教程学神经网络训练,那很可能根本没看过如何收集这些信息用来训练后分析。尤其是“层输出”,因为层通常被包在整个模型里了,你要么得“解包”,要么得修改 layer 的代码来记录输出。
下面是一个简化版的 TensorFlow 自定义训练循环,能收集 loss、metrics、参数、梯度、层输出、层输出梯度的原始值:
import tensorflow as tf
import tqdm
def custom_fit(model, train_dataset, epochs, optimizer):
monitoring_model = tf.keras.Model(
inputs=model.inputs,
outputs=[model.outputs] + [layer.output for layer in model.layers])
loss_history = []
metrics_history = []
variables_history = []
gradients_history = []
layer_outputs_history = []
layer_gradients_history = []
for epoch in tqdm.tqdm(range(epochs)):
for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
with tf.GradientTape() as tape:
monitoring_outputs = monitoring_model(x_batch_train, training=True)
y_pred = monitoring_outputs[0][0]
layer_outputs = monitoring_outputs[1:]
loss = model.compute_loss(x_batch_train, y_batch_train, y_pred, training=True)
loss = optimizer.scale_loss(loss)
all_grads = tape.gradient(loss, model.trainable_weights + layer_outputs)
trainable_grads = all_grads[:len(model.trainable_variables)]
output_grads = all_grads[len(model.trainable_variables):]
optimizer.apply_gradients(zip(trainable_grads, model.trainable_weights))
metrics = model.compute_metrics(x_batch_train, y_batch_train, y_pred)
loss_history.append(loss)
metrics_history.append(metrics)
variables_history.append([tf.identity(v) for v in model.trainable_variables])
gradients_history.append(trainable_grads)
layer_outputs_history.append(layer_outputs)
layer_gradients_history.append(output_grads)
return loss_history, metrics_history, variables_history, gradients_history, layer_outputs_history, layer_gradients_history
上面这段代码能跑,但效率低,也不支持各种数据集和训练超参数写法,也没有处理 epoch 内的值累积。
另一个问题是太死板 —— 每次都收所有数据。你希望的是能“插入”不同的数据收集器。TensorFlow 是用 callback 机制的,你可以定义自己在训练过程中记录数据的逻辑。所以如果你用 TensorFlow,可以像这样做:
📩 如果你想要这段可复用的数据采集代码、指标计算模块,可以评论区留言或私信【Callback源码】,作者整理了一版实用的封装版。
def MyDataCollectionCallback():
def on_epoch_end(variables, gradients, outputs, output_gradients):
...
custom_fit(model, dataset, epochs, callbacks=[MyDataCollectionCallback()])
这是完全可行的。不过完整的代码比较复杂,不适合贴在这里。如果你要看完整、灵活又高效的 TensorFlow 实现,可以看这个 toolkit 里的代码:
- training-instrumentation-toolkit/train_instrumentation.py::fit()
如果你用的是其他平台,也可以参考这个思路改造一下。
数据可视化的采集方式
在任何实用神经网络中,这四种模型视角里的每一项都是一组大型张量,存储它们会占用大量内存。有几个方法可以降低总内存使用量。
首先,你要试验一下是要按每步存(per-step),还是每个 epoch 存一次(per-epoch)。
如果按 epoch 存,你还需要考虑,是只保留每个 epoch 结束时的值,还是在整个 epoch 中做累积。这取决于你想达成什么目标,但下面是一些建议:
- TensorFlow 默认会在一个 epoch 内计算 loss 的平均值,其他指标估计也是这样。这对一些全局性的统计指标(比如模型输出)是个很好的起点。
- 可训练参数在 epoch 结束时就是完成一次完整训练后的最终状态,直接用它们就行了。
- 梯度在每个 mini-batch 更新时波动很大,这是小批量学习的常规现象。很多波动会相互抵消,所以比较好的方式是整个 epoch 内累加梯度。我把它看作模拟“每个 epoch 一次大步”的 SGD。
- 层输出是每个 mini-batch 都会产生的。它们在每步、每样本中都会变,所以统计时应该对每个样本和每个 mini-batch 累加,最后在 epoch 末再计算。比如 Welford 算法就提供了一种高效且数值稳定的方法来累计计算标准差。
另一种降低内存用量的方法是只存统计值。存哪些统计值要看你想从数据中获取什么结果。你很可能需要先收集完整张量,决定要画哪些图,再通过只计算所需统计值的方式做优化。
注意,有些统计量计算起来比其他更耗资源。越耗资源的,就越会影响训练时间。我们在本文中已经用了几个统计值,下面是目前为止用过的指标,再加上一些有用的统计项,并附上它们的计算开销考虑:
- L2 范数,也叫欧几里得范数:在最早讨论权重分布时介绍过。学术界常用,因为它能生成一个代表张量数值大小的标量。L2 范数的基础形式是张量中所有元素的平方和,因此受张量大小影响大,值很容易上千(完整公式在下方)。适合观察单个张量随时间的变化,但不适合比较张量之间的值尺度。L2 范数计算效率还不错。
- 尺寸归一化范数,也叫 RMS:我更喜欢 L2 范数的归一化版本:L2 范数除以张量大小的平方根。它等价于计算张量的 RMS(均方根)。理论上 L2 和 RMS 的效率是一样的。但实际测试里,TensorFlow 对 RMS 的实现(tf.sqrt(tf.reduce_mean(tf.square(tensor))))比 L2(tf.norm(tensor))快 1.5 到 2 倍(在 TensorFlow 2.18,T4 GPU 上测试,张量有几万个元素)。
- 均值和标准差:这是数据统计中常用的一对。它们计算效率高,即使在大张量 per-step 计算时也几乎不会影响训练时间,ML 文献中也常见使用。但必须注意几个问题:第一,权重通常正负非常平衡,均值接近 0。所以均值代表的只是它们偏离 0 的程度,这通常不是你想看的。第二,梯度通常分布集中、尖峰极高、尾部很长。第三,梯度分布两侧往往不对称。所以不仅均值没用,标准差的含义也可疑。
- 取对数后的均值和标准差:如果你真要用均值/标准差,我推荐你去计算张量元素的“数量级”(即取绝对值后的 log),然后在 log 坐标轴上画图。log 尺度更能反映梯度的自然分布(尖峰+长尾)。使用数量级而非原始值也是为了让 log 图成立,也让均值有意义。现在它代表的是平均“尺度”,本质上近似于梯度的标准差。
- 百分位分布:如果你想要最精确的信息,那就计算张量的完整直方图。TensorBoard 里的权重直方图就是这么做的。我的实验发现,算几个跨度从 0 到 100% 的百分位就够了。在这些方法中,这个是最耗算力的。如果你用在模型参数、层输出或梯度上,尤其是 per-step 计算,会显著拖慢训练速度。在 TensorFlow 里可以用 tfp.stats.percentile(tensor, quantiles=[...]) 来算。
以下是一些上面指标的数学公式:
注意:对于梯度来说,均值几乎总是接近0,所以 E[X] 项近似为0,此时标准差就退化成 RMS。
还有些任务你得自定义指标,比如活跃图里用到的“激活率”和“死亡率”。这就是 callback 风格方法的威力 —— 你可以尝试不同的计算方式,然后插进一个通用训练流程中,无需为每个任务都写新的训练循环。
单屏总览
你去查模型排错的教程,最常见的都是讲权重和梯度。现在我希望你意识到,模型不止两个角度,而是四个:
- 模型参数(变量)
- 参数梯度
- 层输出
- 层输出的梯度
我也演示了如何把这些视角用不同方式可视化,每种图都有不同的排错作用。
最后我想讲讲“单屏视图”的价值 —— 把所有这些信息合成一个仪表盘,统一展示整个模型训练过程。
一个模型可能因为很多原因训练失败、训练缓慢或出问题。loss 和 accuracy 图根本不会告诉你背后到底发生了什么。它们只是提示你“可能出问题了”,有时候甚至连这个提示都不给。
所以,把前面讲的所有方法结合起来(还可以加更多),你就能更早发现问题。你怎么组合这些方法,要看你自己在意什么。
我这边用的是下面这种仪表盘:
模型训练总览图
这张图里,每一个小图都表示模型某一方面:
- 左上角:loss 和其他标准指标(每个 epoch)
- 右上角:神经元死亡相关的最严重信号,后续可以扩展显示更多“告警”信息
- 中下方三列小图:四种模型视角的指标(参数、梯度、输出、输出梯度)
- 左列:各层(尺寸归一化)范数的分布,总览模型的值尺度,尤其是随时间变化的趋势。比如值是不是变得异常小或大。
- 中列:各层尺度的 log 相对图,适合判断梯度消失或爆炸。它只显示层之间的相对尺度,不管全局趋势。
- 右列:激活率和死亡率总览。显示的是激活率的平均/最小/最大值分布,还有最严重的死亡率。这个图的理念是“模型只如其最弱的一层”,所以最坏的要画出来。
这个高层视图能让你一眼发现多个问题,很多原本你是完全发现不了的。我用它来判断有哪些地方需要深入看,然后再用每层的细节图去挖。所有这些加起来,就能深入了解模型内部的行为。
总结
在看这篇文章之前,你很可能是靠下面这种基础的 loss 和 metric 曲线来判断模型训练进度的:
现在你可以拿它和“单屏总览”比一比。
拿它和权重、梯度、层输出的逐层图对比一下。
再和神经元激活率、死亡率的图比比。
还可以对比:你在训练时把梯度、层输出等信息变得容易获取之后,能看到哪些以前看不到的东西。
loss 和 accuracy 是模型输出的指标,那是“外在表现”。但模型内部其实有一整座信息宝库在等你挖掘、利用。
我在这篇文章里展示了我认为有用的几种方法,能让你直接看到模型的内部动态。这些只是几个例子,还有很多别的方式,等你自己去尝试、去创造。
我尽量让你探索起来轻松点 —— 讲了一些技术,给了一些真实的例子,还做了个 toolkit 当起点。
所以现在,请去探索、去收获、去狂欢吧。
但先别走……
📩 想获取本文用到的全部图表、工具、notebook 和复现脚本?评论区回复【训练监控】即可。
🧠 作者整理了一份神经网络调试图谱(PDF版),做项目或搞训练监控的朋友可以评论区或者私信【训练图谱】领取。
🔧 技术实战派|AI软硬件一体解决者,芯片→系统→训练→部署,欢迎私信或评论区交流、聊聊项目落地的细节。
最后,这个系列的后续文章,我会告诉你怎么用这些可视化技巧来排查几个模型训练中常见的问题(文章开头有链接)。
源代码
💼 本博文中的这些可视化方法都是作者工程实战中踩坑总结下来的,尤其适合需要 AI 模型从0部署到端侧的开发者,代码和项目说明都已整理,欢迎在评论区交流获取方法。