johnmarktaylor91/torchlens

TorchLens – PyTorchモデルの詳細な検査

機能 – TorchLensは、PyTorchモデルの順伝播（およびオプションで逆伝播）中に流れるすべてのテンソルを記録するPythonライブラリです。完全な計算グラフDAGを構築し、膨大な演算子ごとのメタデータ（形状、型、FLOPs、タイミング、RNG状態など）を付加し、グラフをPDF/HTMLとして可視化できます。このキャプチャは「証明可能に忠実」であり、ライブラリはキャプチャされたDAGをリプレイし、出力が元の実行と一致することを確認し、不一致があればフラグを立てます。

主な機能

• 1行でのトレース：log = tl.trace(model, input) は、活性化値、勾配、メタデータを含む TraceLog オブジェクトを返します。
• 豊富なセレクター – tl.func('relu')、tl.in_module('classifier')、またはカスタムの論理的な組み合わせなどの述語を使用して、保存、停止、または介入する演算子を選択します。
• 順伝播＋逆伝播のキャプチャ – 演算子ごとの勾配が利用可能です。勾配の保存を特定の層に制限できます。
• 受容野 / 投射野分析 – 経験的な勾配マスクとクロスチェックされる幾何学的クエリ。
• 介入 – キャプチャ中に、ゼロアブレーション、スケーリング、ノイズの追加、テンソルの交換などをその場で実行します。
• 可視化 – 自動生成されたPDF、折りたたみ/展開ビュー、インタラクティブなHTML、Netronエクスポート。
• データセット全体の抽出 – データセットをバッチ処理し、活性化値をディスクに書き込みます。
• パフォーマンスへの配慮 – 選択的キャプチャは生の順伝播よりも高速になる場合があります。完全キャプチャは約14倍遅くなりますが、ベンチマーク済みです。
• あらゆるPyTorchモデルで動作 – 静的、動的制御フロー、再帰ループ、transformers、GNNs、拡散モデル、MoEなど。

典型的なワークフロー

import torch, torchvision.models as models, torchlens as tl
model = models.resnet18(weights='IMAGENET1K_V1').eval()
x = torch.rand(1,3,224,224)
log = tl.trace(model, x)          # capture everything
print(log.summary())             # table of layers, FLOPs, params
print(log['conv2d_1_1'].out.shape)  # access a specific activation
log.draw()                       # PDF of the graph

その後、介入、実行の比較、またはトレースのエクスポートが可能です。

インストール – グラフ描画にはGraphvizが必要です。pip install torchlens（CPUのみのコア）でインストールし、HuggingFaceモデルをサポートするために torchlens[hf] などのオプションを追加できます。

対象ユーザー – ニューラルネットワーク内部の詳細な検査、動的アーキテクチャのデバッグ、または分析、プロービング、モデルベースの介入のための体系的な活性化値/勾配抽出を必要とする研究者およびエンジニア。

制限事項 – 逆伝播のキャプチャはPyTorchでのみ機能します。他のバックエンドは、2回目のADパスを介してリーフレベルの勾配を取得します。プラットフォームのサポートは主にLinux/CPUです。macOS/WindowsのCIチェックは限定的です。

---

リンク

• 論文：https://www.nature.com/articles/s41598-023-40807-0
• Model Menagerie (8.5k個のキャプチャ済みモデルカタログ)：https://modelmenagerie.ai
• チュートリアル、パフォーマンスガイド、APIドキュメントはリポジトリの docs/ および notebooks/ にあります。