PyTorchで同期処理を行う

PyTorch Advent Calendar 2018 2日目の記事です。

PyTorchでは*1、GPUの処理は非同期で実施されます。

なので、

t = time.time()
model(var) # GPUでの推論
print(time.time() - t)

としても、非同期で処理が行われるので正しく処理時間は計算できません。

CUDA_LAUNCH_BLOCKING=1 python run.py

のようにCUDA_LAUNCH_BLOCKING=1をつけると同期処理を行うことができます。

参考

• PyTorch デザインノート : CUDA セマンティクス
• Copy tensor from cuda to cpu is too slow - PyTorch Forums

*1:PyTorch以外のライブラリでも同じだと思います