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" のように設定します。
> を含む場合、 ' や " で囲む必要があるのに注意してください。
ログ出力の例
例として、以下のようなスクリプトを実行する際のログを出力することを考えます。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | import torchimport torch.nn as nnclass MyModel(nn.Module): def forward(self, x, y): return x + ymodel = MyModel()model.eval()x = torch.tensor([1, 2, 3], dtype=torch.float32)y = torch.tensor([4, 5, 6], dtype=torch.float32)with torch.no_grad(): torch.onnx.export(model, (x, y), 'model.onnx', input_names=['x', 'y'], output_names=['output']) |
このスクリプトを実行したときの shape_type_inference.cpp のログをすべて出力する場合、 >>shape_type_inference.cpp を PYTORCH_JIT_LOG_LEVEL 環境変数に設定します。
この状態で上記のスクリプトを実行すると、以下のようなログが出力されます。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | [UPDATE shape_type_inference.cpp:2135] ONNX set dynamic input shape.[UPDATE shape_type_inference.cpp:2143] dynamic axes tensor names:[UPDATE shape_type_inference.cpp:1980] Running ONNX shape inference for node: onnx::Constant[DEBUG shape_type_inference.cpp:2016] Original torch graph: graph(%0 : Float(3, strides=[1], requires_grad=0, device=cpu),[DEBUG shape_type_inference.cpp:2016] %1 : Float(3, strides=[1], requires_grad=0, device=cpu)):[DEBUG shape_type_inference.cpp:2016] %2 : Tensor = onnx::Constant[value={1}](), scope: __main__.MyModel::[DEBUG shape_type_inference.cpp:2016] return ()[DEBUG shape_type_inference.cpp:2018] Cloned torch graph to run shape inference: graph():[DEBUG shape_type_inference.cpp:2018] %0 : Tensor = onnx::Constant[value={1}](), scope: __main__.MyModel::[DEBUG shape_type_inference.cpp:2018] return (%0)(途中省略)[DEBUG shape_type_inference.cpp:2126] Torch graph after shape inference:graph(%x : Float(3, strides=[1], requires_grad=0, device=cpu),[DEBUG shape_type_inference.cpp:2126] %y : Float(3, strides=[1], requires_grad=0, device=cpu)):[DEBUG shape_type_inference.cpp:2126] %output : Float(3, strides=[1], requires_grad=0, device=cpu) = onnx::Add(%x, %y), scope: __main__.MyModel:: # C:\work\git_repos\deform_conv2d_onnx_exporter\sample.py:7:0[DEBUG shape_type_inference.cpp:2126] return (%output)============== Diagnostic Run torch.onnx.export version 2.0.0+cpu ==============verbose: False, log level: Level.ERROR======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ======================== |
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_LEVEL に shape_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 に変更することで、ログを出力できるようになります。
なお、このプルリクエストで修正されました！
このログを使ってデバッグした話の詳細については、次回以降に書くかもしれません。
