1. 程式人生 > 其它 >pytorch: cudnn.benchmark=True

pytorch: cudnn.benchmark=True

技術標籤:torch

import torch.backends.cudnn as cudnn
cudnn.benchmark = True

設定這個 flag 可以讓內建的 cuDNN 的 auto-tuner 自動尋找最適合當前配置的高效演算法,來達到優化執行效率的問題。
如果網路的輸入資料維度或型別上變化不大,也就是每次訓練的影象尺寸都是一樣的時候,設定 torch.backends.cudnn.benchmark = True 可以增加執行效率;
如果網路的輸入資料在每次 iteration 都變化的話,會導致 cnDNN 每次都會去尋找一遍最優配置,這樣反而會降低執行效率

參考文獻:

BOKEYUAN