TrainTools

traintools

PyTorch 向けの、論文に基づく ML トレーニング診断ツールです。実行中に実用的な疑問に答える、小さなツール群です。

GitHub のソース | PyPI

実際の実行結果や誤検知を共有したい場合は、 診断レポートのテンプレートをご利用ください。

pip install traintools

問題に適したツールがわからない場合は、こちらをご覧ください。

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

エージェントは、JSON CLI、エージェントガイド、llms.txt、 またはオプションのローカル MCP サーバーを利用できます。

pip install "traintools[mcp]"
traintools-mcp

同じローカルの読み取り専用サーバーは、公式 MCP Registry からも見つけられ、永続的にインストールせずに起動できます。

uvx --with mcp traintools mcp

AI コーディングアシスタントの場合：まず traintools recommend "<training symptom>" --json を実行し、ユーザーが既に使用しているループに最小限の診断スニペットを挿入します。エージェント向けのエントリーポイントは、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()])

Gradient Noise Scale

GNS は、サンプルごとの勾配分散と勾配シグナルの比率です。

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

これは臨界バッチサイズ B* を推定します。これを超えると、バッチを大きくしても最適化の進展があまり得られなくなります。

• GNS > B：バッチサイズが小さすぎます。バッチを大きくすると改善する可能性があります
• GNS < B：バッチサイズが大きすぎます。必要以上に大きい可能性があります
• GNS ~= B：効率的な最適領域に近い状態です

traintools は McCandlish et al. 2018 による不偏推定量（ベッセル補正済み分散、バイアス補正済みシグナル）を使用し、GNS を EMA(tr(Sigma)) / EMA(||G||^2) として追跡します。

[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 は、あまり活用されていない次の2つの学習ダイナミクスのプローブを実装します。

• Toneva et al. 2019 によるサンプルの忘却イベント
• Swayamdipta et al. 2020 によるデータセット・カルトグラフィー形式の信頼度と変動性

通常の分類実行中に、安定したデータセット ID、ロジット、ラベルを使用してください。

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 et al. 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 et al. 2020 による Area Under the Margin 統計量を実装します。各サンプルについて、次の値を平均します。

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 et al. 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, and Donoho 2020 による学習後期の分類器の幾何学的特性を測定します。

• NC1：クラス内の特徴の変動性とクラス間の変動性の比率
• NC2：クラス平均が単体 ETF の幾何学的特性に近づいているか
• NC3：分類器の重みがクラス平均と整列しているか（重みが指定された場合）
• NCC accuracy：最近傍クラス中心の精度

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 はアルファ版のソフトウェアです。診断機能は意図的に小さく、十分にテストされていますが、しきい値はヒューリスティックであり、自動的な正解ではなく、学習の指標として解釈してください。バグ報告、ベンチマークのトレース、実際の失敗事例を特に歓迎します。

問題別のガイドはドキュメントにあります。診断オブジェクトは 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