SyntaxHighlighter

2023年4月3日月曜日

PyTorch の TorchScript のデバッグメッセージ (PYTORCH_JIT_LOG_LEVEL によるログ出力)

PyTorch の挙動を調べる必要があった際に、ログ出力の設定を調べたことを書いておこうと思います。

PyTorch の TorchScript のログ出力

PyTorch のモデルを保存する際などに PyTorch 内部で利用される TorchScript には、標準でデバッグ用のログ出力の機能が備わっています。

なお、私は ONNX へのエクスポートの際の挙動のデバッグの際にこのログ出力の機能を用いました。

PYTORCH_JIT_LOG_LEVEL 環境変数の設定

基本的に、 jit_log.h の冒頭のコメント部分に書かれていることに従って環境変数を設定します。

この環境変数には、以下のようにログ出力したいファイルのファイル名を指定します。

  • PyTorch のソースコードを見て、ログを出力したいファイルを探します。
    例えば、 torch/csrc/jit/passes/onnx/shape_type_inference.cpp のログを出力したい場合、 shape_type_inference.cpp が設定すべきファイル名になります。
  • 出力したいログのレベルは現時点では 3段階あり、その値に応じて > をいくつか付与します。
  • 複数のファイルのログを出力したい場合は、 : で区切って連結します。
  • 結果として、 shape_type_inference.cpp:>subgraph_matcher.cpp のような文字列を PYTORCH_JIT_LOG_LEVEL に設定します。

結果として、 Bash であれば export 'PYTORCH_JIT_LOG_LEVEL=shape_type_inference.cpp:>subgraph_matcher.cpp' のように設定します。
Windows 上のコマンドプロンプトを使う場合は、 set "PYTORCH_JIT_LOG_LEVEL=shape_type_inference.cpp:>subgraph_matcher.cpp" のように設定します。

> を含む場合、 '" で囲む必要があるのに注意してください。

ログ出力の例

例として、以下のようなスクリプトを実行する際のログを出力することを考えます。

このスクリプトを実行したときの shape_type_inference.cpp のログをすべて出力する場合、 >>shape_type_inference.cppPYTORCH_JIT_LOG_LEVEL 環境変数に設定します。
この状態で上記のスクリプトを実行すると、以下のようなログが出力されます。

Windows 上でログ出力する際の注意事項

実際には、先の例の通りにやっても、 Windows ではログが出力されません。
これは、 v2.0.0 の時点での (おそらく) バグで、 PYTORCH_JIT_LOG_LEVEL に設定されたファイル名と __FILE__ を比較する際の不具合によるものです。

StripBasename 関数の中でファイルのフルパスからファイル名部分を抽出していますが、この部分はディレクトリの区切り文字として / しか見ていません。
そのため、 Windows で __FILE__C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\torch\csrc\jit\passes\onnx\shape_type_inference.cpp のようなファイル名が設定されている際には、ファイル名部分として shape_type_inference.cpp がうまく抽出されず、 PYTORCH_JIT_LOG_LEVELshape_type_inference.cpp が設定されていてもうまくログが出力されませんでした。

この場合、 StripBasename を修正してビルドし直すか、 site-packages\torch\lib\torch_python.dll をバイナリエディタで直接修正し、 DLL に含まれる C:\actions-runner\_work\pytorch\pytorch\builder\windows\pytorch\torch\csrc\jit\passes\onnx\shape_type_inference.cpp の部分を C:/actions-runner/_work/pytorch/pytorch/builder/windows/pytorch/torch/csrc/jit/passes/onnx/shape_type_inference.cpp に変更することで、ログを出力できるようになります。

なお、このプルリクエストで修正されました!

このログを使ってデバッグした話の詳細については、次回以降に書くかもしれません。