TrainTools

traintools

基于论文的 PyTorch 训练诊断工具。用小巧的工具，在训练仍在运行时解答实际问题。

GitHub 上的源代码 | PyPI

有真实的训练过程或误报案例可以分享吗？请使用 诊断报告模板。

pip install traintools

不确定哪个工具适合当前问题？

traintools recommend "my loss became NaN and gradients explode"
traintools integration gradient-health --framework pytorch

Agent 可以使用 JSON CLI、Agent 指南、llms.txt， 或可选的本地 MCP 服务：

pip install "traintools[mcp]"
traintools-mcp

同一个本地只读服务也可通过官方 MCP Registry 发现，并且无需持久安装即可启动：

uvx --with mcp traintools mcp

面向 AI 编码助手：先运行 traintools recommend "<training symptom>" --json，然后将最小的诊断代码片段插入用户现有的循环中。面向 Agent 的入口包括 AGENTS.md、 llms.txt、JSON 架构以及可选的 MCP 服务。

工具

工具	解答的问题
Gradient Noise Scale (GNS)	我的批次大小是否在浪费计算资源？
GradientAccumulationGNS	我能否在梯度累积期间免费获得 GNS？
PlasticityProbe	我的网络是否正在失去学习能力？
TrainGuard	现在应该停止训练吗？
BatchInspector	这个批次是否存在错误、失衡或尺度不当？
GradientHealthMonitor	梯度是否有限、被裁剪、消失、爆炸，或相对于权重过大？
ExampleDynamicsTracker	哪些样本被遗忘、难以学习、存在歧义或可能被错误标注？
GradientConfusionMonitor	微批次梯度是否相互冲突并拖慢 SGD？
AUMTracker	根据边际动态，哪些样本看起来被错误标注？
EL2NTracker	训练早期哪些样本很重要或可以剪枝？
NeuralCollapseMonitor	分类器是否进入了神经坍缩几何状态？

快速开始

from traintools import BatchInspector, GradientHealthMonitor
from traintools.callbacks.pytorch import TraintoolsTracker

tracker = TraintoolsTracker(model, loss_fn)
batch_inspector = BatchInspector(expected_num_classes=10)
grad_health = GradientHealthMonitor(max_grad_norm=1.0)

for step, (x, y) in enumerate(dataloader):
    batch_report = batch_inspector.inspect(x, y, step=step)
    if not batch_report.ok:
        print(batch_report)

    loss = loss_fn(model(x), y)
    loss.backward()

    grad_report = grad_health.inspect(model, step=step, lr=optimizer.param_groups[0]["lr"])
    if not grad_report.ok:
        print(grad_report)

    optimizer.step()
    optimizer.zero_grad()

    decision = tracker.step(step=step, inputs=x, targets=y, val_loss=val_loss)
    if decision and decision.should_stop:
        break

HuggingFace Trainer：

from transformers import Trainer
from traintools.callbacks.huggingface import TraintoolsCallback

trainer = Trainer(model=model, ..., callbacks=[TraintoolsCallback()])

梯度噪声尺度

GNS 是逐样本梯度方差与梯度信号的比值：

GNS = tr(Sigma) / ||G||^2

它用于估算临界批次大小 B*：超过这个点后，增大批次带来的优化进展会变得有限。

• GNS > B：批次过小，增大批次可能会有帮助
• GNS < B：批次过大，当前批次可能超出所需大小
• GNS ~= B：接近高效前沿

traintools 使用 McCandlish 等人 2018 年提出的无偏估计量 （经贝塞尔校正的方差、经偏差校正的信号），并按 EMA(tr(Sigma)) / EMA(||G||^2) 跟踪 GNS。

[step 500] GNS=5010.7 (EMA)  critical_batch=5011  current=64  regime=under-batched
  > Batch size 64 is ~78x below the critical batch (~5011). Larger batches would give cleaner gradients per step.

梯度累积期间免费计算 GNS

如果你已经在使用梯度累积，那么你本来就要计算的逐微批次梯度，正是 GNS 所需的样本。

from traintools import GradientAccumulationGNS

gns = GradientAccumulationGNS(model, micro_batch_size=B_micro)

for step in range(num_steps):
    for xm, ym in micro_batches:
        (loss_fn(model(xm), ym) / accum_steps).backward()
        gns.record_microbatch()
    optimizer.step()
    result = gns.compute(step=step)
    optimizer.zero_grad()
    gns.reset_accumulation()

PlasticityProbe

PlasticityProbe 直接测量激活：

• 休眠单元比例：对每个输入，激活都接近零的单元所占比例
• 特征有效秩：激活协方差的归一化有效秩

这些指标会合并为一个取值范围为 [0, 1] 的可塑性分数。

[step 200] Plasticity Score: 0.706
  All layers healthy.

TrainGuard

TrainGuard 会对验证损失拟合幂律或指数曲线，通过自举估算不确定性，并且只有在继续训练看起来不太可能带来显著效果时才会停止。

[step 400] STOP
  current loss: 0.6536
  predicted final: 0.6119
  expected improvement: 0.0417 (90% CI: [0.0012, 0.0821])
  estimated plateau at step: 3200
  reason: No improvement in 300 steps (best=0.6350 at step 93).

BatchInspector

BatchInspector 会捕获异常张量和标签，避免它们悄然污染一次运行。

from traintools import BatchInspector

inspector = BatchInspector(expected_num_classes=10, max_abs_value=1e4)
report = inspector.inspect(inputs=x, targets=y, step=step)
if not report.ok:
    print(report)

它会检查空张量、NaN/inf、极端尺度、常量张量、超出预期类别范围的标签，以及严重的批次不平衡。

GradientHealthMonitor

GradientHealthMonitor 在 backward() 之后、optimizer.step() 之前调用。

from traintools import GradientHealthMonitor

monitor = GradientHealthMonitor(max_grad_norm=1.0)
loss.backward()
report = monitor.inspect(model, step=step, lr=optimizer.param_groups[0]["lr"])
if not report.ok:
    print(report)

它会报告全局和逐层梯度范数、非有限梯度、可能的梯度消失/爆炸、裁剪系数，以及更新量与权重的比率。

ExampleDynamicsTracker

ExampleDynamicsTracker 实现了两种尚未充分利用的训练动态探测方法：

• 来自 Toneva 等人（2019）的样例遗忘事件
• 来自 Swayamdipta 等人（2020）的数据集制图式置信度和可变性

在常规分类训练过程中，使用稳定的数据集 ID、logits 和标签。

from traintools import ExampleDynamicsTracker

dynamics = ExampleDynamicsTracker()

for step, (ids, x, y) in enumerate(dataloader):
    logits = model(x)
    loss = loss_fn(logits, y)
    dynamics.update(ids, logits, y, step=step)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

print(dynamics.summary())
print("likely noisy or brittle:", [ex.example_id for ex in dynamics.most_forgotten(20)])
print("ambiguous:", [ex.example_id for ex in dynamics.cartography_region("ambiguous")])

遗忘事件是指同一样例从分类正确转变为分类错误。反复被遗忘的样例往往存在歧义、标签错误，或处于分布边缘。永不被遗忘的样例可作为剪枝或课程学习实验的候选对象。

GradientConfusionMonitor

GradientConfusionMonitor 遵循 Sankararaman 等人（2019）提出的梯度冲突概念，用于估算微批次梯度是相互一致还是彼此冲突。

from traintools import GradientConfusionMonitor

confusion = GradientConfusionMonitor(n_splits=4)
report = confusion.estimate(model, loss_fn, x, y, step=step)
if not report.ok:
    print(report)

它会报告成对梯度余弦值的均值/最小值/最大值、负梯度对的比例，以及简洁的冲突分数。冲突较高可能意味着标签噪声、样本不兼容、深度/初始化问题，或需要采用不同的批次/课程学习策略。

AUMTracker

AUMTracker 实现了 Pleiss 等人（2020）提出的间隔下面积统计量。对于每个样例，它会计算以下值的平均值：

true_class_logit - max(other_class_logits)

AUM 较低的样例可作为标签审查或歧义检查的候选对象。

from traintools import AUMTracker

aum = AUMTracker(low_aum_threshold=0.0)
for step, (ids, x, y) in enumerate(dataloader):
    logits = model(x)
    aum.update(ids, logits, y, step=step)

print([ex.example_id for ex in aum.lowest_aum(20)])

EL2NTracker

EL2NTracker 实现了 Paul 等人（2021）提出的低成本样例重要性评分：

||softmax(logits) - one_hot(label)||_2

EL2N 较高的样例往往较重要、较难、带有噪声，或处于分布边缘。EL2N 较低的样例可作为数据剪枝实验的候选对象。

from traintools import EL2NTracker

el2n = EL2NTracker()
el2n.update(ids, logits, y, step=step)
important = el2n.highest(100)
prune_candidates = el2n.lowest(100)

NeuralCollapseMonitor

NeuralCollapseMonitor 用于衡量 Papyan、Han 和 Donoho（2020）提出的训练后期分类器几何特性：

• NC1：类内特征可变性相对于类间可变性的比例
• NC2：类别均值趋近单纯形 ETF 几何结构的程度
• NC3：分类器权重与类别均值的对齐程度（提供权重时）
• NCC 准确率：最近类别中心准确率

from traintools import NeuralCollapseMonitor

collapse = NeuralCollapseMonitor()
report = collapse.measure(features, labels, classifier_weight=model.fc.weight)
print(report)

安装

# Core: PyTorch only
pip install traintools

# Curve fitting and plotting helpers
pip install traintools[full]

# HuggingFace Trainer integration
pip install traintools[hf]

# Local MCP server for compatible AI clients
pip install traintools[mcp]

# Development
pip install -e ".[dev]"

项目状态

traintools 是 alpha 阶段的软件。诊断工具有意保持轻量且经过充分测试，但阈值只是启发式指标，应将其视为训练信号，而非自动判定的事实。我们尤其欢迎 Bug 报告、基准测试记录和真实故障案例。

面向问题的指南请参阅文档。可以使用 write_json_report 将诊断对象写入带版本信息的 JSON 文件。

参考文献

• McCandlish、Kaplan、Amodei 等人，2018 年。大批量训练的经验模型。
• Dohare、Sutton 等人，2024 年。深度持续学习中的可塑性丧失。
• Lyle 等人，2023 年。理解神经网络中的可塑性。
• Domhan 等人，2015 年。通过外推学习曲线加速 DNN 的自动超参数优化。
• Toneva 等人，2019 年。深度神经网络学习中样例遗忘的经验研究。
• Swayamdipta 等人，2020 年。数据集制图：利用训练动态映射和诊断数据集。
• Sankararaman 等人，2019 年。神经网络过参数化对梯度混淆和随机梯度下降的影响。
• Pleiss 等人，2020 年。使用边际排序曲线下面积识别错误标记的数据。
• Paul、Ganguli 和 Dziugaite，2021 年。精简数据的深度学习：在训练早期寻找重要样例。
• Papyan、Han 和 Donoho，2020 年。深度学习训练终末阶段神经坍缩的普遍性。

许可证

MIT