PyTorchでプロファイルを取得する。torch.profiler.profile
torch.profiler.profileの機能を使ってプロファイルを取ることができる。プロファイルとは要するにどの関数や処理でどれほどの時間を要しているのかを計測すること。
計測
# profiler setting if config.profile: prof = torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA ], with_stack=True, with_flops=True # FLOPs推定も含める ) prof.start() else: prof = None
のような感じでプロファイル用のオブジェクトを作成して、計測したいところをprof.start()とprof.stop()で囲めば良い。例えば
prof.start() loss.backward() prof.stop()
みたいな感じ。
結果の出力
結果の出力はテーブル形式やファイル出力ができる。
テーブル形式での表示
テーブル形式での表示は、次のように書くと、プロファイルをとった部分について総cpu時間でソートして、上位10行だけ表示する。
Nameの項を省略させたくない場合はmax_name_column_widthで大きめの値を指定してやれば良い。
print(prof.key_averages().table( sort_by="self_cpu_time_total", row_limit=10) )
出力例は
--------------------------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls Total MFLOPs
--------------------------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
autograd::engine::evaluate_function: torch::autograd::AccumulateGrad 1.28% 1.106ms 56.98% 49.355ms 9.871ms 0.000us 0.00% 82.430us 16.486us 5 --
FullyShardedDataParallel._post_backward_hook 49.58% 42.939ms 55.40% 47.978ms 9.596ms 0.000us 0.00% 82.430us 16.486us 5 --
mpi:_reduce_scatter_base 0.00% 0.000us 0 35.221ms 7.044ms 0.000us 0.00% 0.000us 0.000us 5 --
autograd::engine::evaluate_function: BackwardHookFunctionBackward 1.91% 1.658ms 30.33% 26.267ms 2.919ms 0.000us 0.00% 0.000us 0.000us 9 --
FullyShardedDataParallel._pre_backward_hook 4.56% 3.950ms 28.17% 24.400ms 4.067ms 0.000us 0.00% 0.000us 0.000us 6 --
FullyShardedDataParallel._pre_backward_prefetch 22.69% 19.650ms 23.38% 20.253ms 4.051ms 0.000us 0.00% 0.000us 0.000us 5 --
mpi:_allgather_base 0.00% 0.000us 0 16.053ms 3.211ms 0.000us 0.00% 0.000us 0.000us 5 --
autograd::engine::evaluate_function: AddmmBackward0 0.95% 825.352us 4.84% 4.190ms 837.935us 0.000us 0.00% 117.728us 23.546us 5 --
AddmmBackward0 0.84% 727.863us 3.05% 2.645ms 529.074us 0.000us 0.00% 92.000us 18.400us 5 --
aten::fill_ 0.22% 193.891us 2.52% 2.182ms 436.461us 23.551us 0.41% 25.727us 5.145us 5 --
aten::ones_like 0.06% 53.882us 2.44% 2.114ms 2.114ms 0.000us 0.00% 4.352us 4.352us 1 --
Activity Buffer Request 2.19% 1.898ms 2.19% 1.898ms 1.898ms 2.176us 0.04% 2.176us 2.176us 1 --
aten::chunk 0.12% 100.948us 1.48% 1.286ms 257.146us 0.000us 0.00% 0.000us 0.000us 5 --
aten::t 0.62% 538.507us 1.48% 1.285ms 67.654us 0.000us 0.00% 0.000us 0.000us 19 --
autograd::engine::evaluate_function: ViewBackward0 0.59% 509.453us 1.43% 1.236ms 123.593us 0.000us 0.00% 0.000us 0.000us 10 --
autograd::engine::evaluate_function: SplitWithSizesBackward0 0.45% 388.952us 1.40% 1.210ms 242.022us 0.000us 0.00% 64.960us 12.992us 5 --
aten::split 0.38% 328.214us 1.37% 1.185ms 236.956us 0.000us 0.00% 0.000us 0.000us 5 --
aten::pad 0.09% 76.219us 1.24% 1.075ms 358.439us 0.000us 0.00% 48.191us 16.064us 3 --
aten::slice 0.93% 803.974us 1.21% 1.046ms 45.471us 0.000us 0.00% 0.000us 0.000us 23 --
aten::narrow 0.40% 350.577us 1.21% 1.046ms 58.086us 0.000us 0.00% 0.000us 0.000us 18 --
aten::constant_pad_nd 0.38% 330.485us 1.15% 999.097us 333.032us 0.000us 0.00% 48.191us 16.064us 3 --
aten::mm 0.93% 802.129us 1.11% 962.021us 106.891us 92.000us 1.60% 92.000us 10.222us 9 23.110
autograd::engine::evaluate_function: ReluBackward0 0.33% 284.475us 1.05% 910.788us 227.697us 0.000us 0.00% 8.640us 2.160us 4 --
FullyShardedDataParallel._post_backward_prefetch 1.01% 874.844us 1.01% 874.844us 174.969us 0.000us 0.00% 0.000us 0.000us 5 --
SplitWithSizesBackward0 0.18% 159.000us 0.93% 806.363us 161.273us 0.000us 0.00% 64.960us 12.992us 5 --
aten::div_ 0.73% 629.432us 0.90% 781.302us 78.130us 34.239us 0.60% 34.239us 3.424us 10 --
autograd::engine::evaluate_function: TBackward0 0.39% 338.954us 0.89% 772.324us 154.465us 0.000us 0.00% 0.000us 0.000us 5 --
aten::transpose 0.61% 530.956us 0.86% 746.928us 39.312us 0.000us 0.00% 0.000us 0.000us 19 --
ViewBackward0 0.30% 263.649us 0.84% 726.474us 72.647us 0.000us 0.00% 0.000us 0.000us 10 --
cudaLaunchKernel 0.76% 655.898us 0.76% 655.898us 15.617us 0.000us 0.00% 0.000us 0.000us 42 --
--------------------------------------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
chrome trace形式での出力
プロファイルした項目の依存関係とか、処理の時間推移を可視化したい場合にこの方法を使う。
prof.export_chrome_trace("trace.json")
でjson形式でプロファイル結果を出力して、chrome://tracing/ にアクセスし、このファイルを読み込むと下の画像のように整形されて表示される。横軸は時間。
