モデルの推論・学習時間を評価する（PyTorch, profiler）
なんか学習時間長くないか？？？
って思う瞬間、ありますよね。私は研究を進めていて、結構この悩みがあります。学習待ってる時間って虚無過ぎませんか？
どこが長くなっているのか特定するために、start = time.time()とend = time.time()で処理を挟んで、print(end - start)って書いて、処理ごとに同じことを繰り返して…
やってられません。しかも、これはCPUの実行待ちの間もどんどん計測されてしまって精度もよくないし、得られる情報が少なすぎます。
そこで、今回はPyTorchに用意されているtorch.profilerを使って詳細にモデルのボトルネックを特定してみます。
まずはインポート
まず、必要なクラスや関数をインポートしていきます。
from torch.profiler import profile, record_function, ProfilerActivity
-
profileは、処理したい部分全体のコードブロックをwith文で指定して実行時間を計測します。 -
record_functionは、profileで指定したコードブロックの中で、任意の処理部分を同様にwith文で名前をつけて計測項目に加えます。 -
ProfilerActivityは、CPUによる処理やGPUによる処理を指定して計測するための変数クラスです。ProfilerActivity.CPUやProfilerActivity.CUDAのように指定します。
profilerの使い方
あくまで一例ですが、一般的な学習ループでの使用法を示します。
例えば、バッチ処理部分で各処理にかかる時間や資源を調査するには以下のようにします。
with profile(activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
for j, (x, t) in enumerate(train_loader):
optimizer.zero_grad()
x, t = x.to(device), t.to(device)
with record_function("inference"):
y = model(x)
loss = criterion(y, t)
with record_function("backward"):
loss.backward()
with record_function("step"):
optimizer.step()
current_loss += loss.detach()
print(prof.key_averages().table(sort_by = "cuda_time", row_limit = 10))
with profile(activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]) as profによって、処理のうちCPUとCUDAを使用したもののレコードを保存します。
record_functionをwith文で使用して計測したい処理部分に名前を付けます。
record_functionを使用しなくてもtorchに関連したすべての処理の詳細なレポートを確認できますが、サンプルのように推論全体にかかる時間や逆伝播全体にかかる時間などを計測したい場合はこれらの処理をrecord_functionで指定したほうがわかりやすいと思います。
該当の処理が終了したら、print(prof.key_averages().table(sort_by = "cuda_time", row_limit = 10))で詳細なレポートを表示します。
key_averages()で重複するイベントごとに実行時間などの平均をとり、table(sort_by = "cuda_time", row_limit = 10)で1プロセス当たりのGPU使用時間で降順にソート、上位10個までを表示します。
最後にまとめて表示してくれるし、処理ごとに明示的にラベリングできるし、プロセスごとの実行時間だけでなくトータルの時間やメモリ使用率、処理が呼び出された回数などもわかります。
なんて便利なんでしょう。
実行結果がこちらになります。
------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Name Self CPU % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg # of Calls
------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
autograd::engine::evaluate_function: Con... 0.01% 3.414ms 0.06% 35.799ms 127.854us 2.134ms 0.00% 35.657s 127.347ms 280
ConvolutionBackward0 0.01% 3.795ms 0.05% 32.286ms 115.307us 1.233ms 0.00% 35.655s 127.339ms 280
aten::convolution_backward 0.04% 23.487ms 0.04% 28.490ms 101.752us 35.509s 54.31% 35.654s 127.335ms 280
aten::_foreach_div 4.95% 3.143s 4.96% 3.145s 112.327ms 2.848s 4.36% 2.864s 102.281ms 28
aten::_foreach_add_ 0.04% 27.554ms 0.04% 27.705ms 494.738us 5.043s 7.71% 5.062s 90.402ms 56
aten::_foreach_div_ 0.00% 1.675ms 0.00% 1.675ms 59.819us 2.352s 3.60% 2.352s 83.984ms 28
autograd::engine::evaluate_function: Cud... 0.00% 2.527ms 0.02% 15.586ms 79.520us 556.000us 0.00% 6.483s 33.077ms 196
CudnnBatchNormBackward0 0.00% 2.037ms 0.02% 13.059ms 66.630us 570.000us 0.00% 6.483s 33.074ms 196
aten::cudnn_batch_norm_backward 0.02% 9.998ms 0.02% 11.022ms 56.237us 6.473s 9.90% 6.482s 33.071ms 196
enumerate(DataLoader)#_MultiProcessingDa... 1.83% 1.164s 1.85% 1.174s 40.495ms 763.384ms 1.17% 765.128ms 26.384ms 29
------------------------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 63.455s
Self CUDA time total: 65.381s
各列はそれぞれ以下の内容を表しています。
-
Self DEVICE %：処理時間全体に占める、該当処理の、サブプロセスを含まない実行時間の割合 -
Self DEVICE：該当処理の、サブプロセスを含まない実行時間 -
DEVICE total %：処理時間全体に占める、該当処理の、サブプロセスを含む実行時間の割合（コードブロックの実行が始まって終わるまでの時間） -
DEVICE total：該当処理の、サブプロセスを含む実行時間 -
DEVICE time avg：該当処理1つあたりの実行時間（例えばconv1層分の時間など） -
# of Calls：該当処理の実行回数
DEVICEをCPUやGPUと読み替えてください。サブプロセス云々の話は、例えば
with record_function("process A"):
hogehogeIO(fuga)
processB(piyo)
def processB(piyo):
with record_function("process B"):
toto
titi
...
といったときを考えます。
hogehogeIOの入出力待ち時間などがあったり、processBの処理がプロファイリングされているとき、
Self DEVICEはこれらの待ち時間やプロファイルされている処理時間を除外した処理時間を出力し、DEVICE totalはprocessBの実行時間も含めて出力する、ということだと思います。
上位3つのプロセスを見ると、すべてConvolutionのBackward関係になっており、
CUDA totalがほぼ同じで、Self CUDA %がそれぞれ0, 0, 54 %になっています。
autograd::engine...とConvolutionBackward0、aten::convolution_backwardが順に
呼び出され、順に処理待ちの状態になっている関係がわかると思います。
profileで学習のボトルネックを効率的に見つけよう
これらの実験によって、畳み込み層の逆伝播が学習のボトルネックになっていることがわかりました。
あまり層を深くしたつもりもなく、てっきり学習ループのコーディングや損失の設計がミスっているのかと思っていましたが単純にモデルが重いだけでした。
ゆったり学習して成長を見守りたいと思います。
Discussion