Stonesjtu/pytorch_memlab

pytorch_memlab – PyTorch用の軽量CUDAメモリデバッグツールボックス

何ができるか – Pythonだけで動作するパッケージで、PyTorchコードの各行がどのくらいのGPUメモリを消費しているかを可視化し、デバイス上に存在するテンソルを一覧表示し、一時的にすべてをCPUに移動できます。コマンドラインでも、Jupyter/IPythonのマジックコマンド経由でも動作します。

主な機能

行単位のメモリプロファイラー (LineProfiler / @profileデコレータ) – line_profilerに似ていますが、各ソース行ごとにCUDAのactive_bytesとreserved_bytesを報告します。
メモリレポーター (MemReporter) – すべてのライブtorch.Tensorオブジェクトを列挙し、下位のUntypedStorageのサイズを表示できます。共有重み、勾配、詳細なストレージ再利用図も対応しています。
カジュアルモード – 必要に応じてすべてのGPUテンソルをCPUに戻すヘルパークラス。チェックポイントを保存せずにGPUを解放するのに便利です。
IPython統合 – %mlrunと%%mlrunマジックコマンドで、1つのコマンドで関数やノートブックセル全体をプロファイリングできます。
GPU選択 – set_target_gpuで、実行中にプロファイラーが監視するデバイスを切り替えることができます。

インストール方法

# PyPIからの安定版
pip install pytorch_memlab

# GitHubから最新コードを直接インストール
pip install git+https://github.com/stonesjtu/pytorch_memlab

(IPythonマジックを使用するにはオプションの追加パッケージが必要です: pip install pytorch_memlab[ipython]。)

一般的なワークフロー

関数のプロファイリング
```
from pytorch_memlab import profile

@profile
def train_step(x):
    net = torch.nn.Linear(1024, 1024).cuda()
    return net(x).mean()
```
スクリプト終了後に、各行がどのくらいのGPUメモリを割り当てたかを示すテーブルが出力されます。

ライブテンソルの確認

from pytorch_memlab import MemReporter
reporter = MemReporter(model)   # modelはオプション
reporter.report()               # 簡潔なテーブルを出力
reporter.report(verbose=True)  # ストレージ共有の矢印を表示

ノートブックでの使用

%load_ext pytorch_memlab
%%mlrun -f train_step
train_step(torch.randn(512, 1024).cuda())

GPUメモリの一時解放

from pytorch_memlab import Courtesy
c = Courtesy()
c.yield_memory()   # テンソルをCPUに移動
# …信号を待つ…
c.restore()        # テンソルを戻す

なぜ重要か – 深層学習モデルの開発では、メモリ不足（OOM）によるクラッシュがよく発生します。pytorch_memlabは隠れたCUDAの割り当て動作を可視化し、非効率なテンソルのライフタイム、意図しないストレージ共有、または忘れられた中間バッファを特定するのに役立ちます。

制限事項 / 既知の問題

プロファイラーはPythonが認識しているテンソルしか見ることができません。autogradによって作成された低レベルのC側バッファは「見えない」として報告され、それでもメモリを消費します。
すべてのテンソルがCPU上にあっても、PyTorchのCUDAコンテキスト自体が約1GBのGPUメモリを確保しており、現在のところこのメモリはライブラリで解放できません。
CUDA対応GPUでのみ動作します。AMDやCPUオンリーの実行には対応していません。

プロジェクトの状態 – 活発にメンテナンス中（最新リリース 0.3.2、2026年6月）。CIでテストを実行し、PyPIにウェHEELを公開しています。READMEには使用例、デモノートブック、変更履歴が含まれています。

上記のすべての情報はリポジトリのREADMEから直接取得したものであり、外部からの主張は一切追加されていません。

Stonesjtu/pytorch_memlab

pytorch_memlab – PyTorch用の軽量CUDAメモリデバッグツールボックス

関連