快捷導(dǎo)航

Pytorch精準(zhǔn)記錄函數(shù)運行時間的方法

更新時間：2024年11月11日 09:23:35 作者：非線性光學(xué)元件

參考Pytorch官方文檔對CUDA的描述,GPU的運算是異步執(zhí)行的,一般來說,異步計算的效果對于調(diào)用者來說是不可見的,異步計算的后果是,沒有同步的時間測量是不準(zhǔn)確的,所以本文給大家介紹了Pytorch如何精準(zhǔn)記錄函數(shù)運行時間,需要的朋友可以參考下

0. 引言

參考Pytorch官方文檔對CUDA的描述，GPU的運算是異步執(zhí)行的。一般來說，異步計算的效果對于調(diào)用者來說是不可見的，因為

每個設(shè)備按照排隊的順序執(zhí)行操作
Pytorch對于CPU和GPU的同步，GPU間的同步是自動執(zhí)行的，不需要顯示寫在代碼中

異步計算的后果是，沒有同步的時間測量是不準(zhǔn)確的。

1. 解決方案

參考引言中提到的幫助文檔，Pytorch官方給出的解決方案是使用torch.cuda.Event記錄時間，具體代碼如下：

# import torch
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()

# Run your code snippet here

end_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)  # elapsed time (ms)

將你的代碼插入start_event.record()和end_event.record()中間以測量時間（單位毫秒）。

有能力的讀者也可以包裝為裝飾器或者with語句使用：

先書寫一個自定義with類（ContextManager）

class CudaTimer:
    def __init__(self):
        self.start_event = torch.cuda.Event(enable_timing=True)
        self.end_event = torch.cuda.Event(enable_timing=True)

    def __enter__(self):
        self.start_event.record()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_event.record()
        torch.cuda.synchronize()
        self.elapsed_time = self.start_event.elapsed_time(self.end_event) / 1000 # ms -> s

再安裝如下with語句返回：

with CudaTimer() as timer:
	# run your code here
dt = timer.elapsed_time  # s

這樣保證了多個文件調(diào)用時語句的簡單性。特別提醒：獲取timer.elapsed_time操作不要寫在with語句內(nèi)部。在with語句未結(jié)束時，是無法獲取timer的成員變量的。

2. 補充

對于CPU和GPU混合操作的函數(shù)，使用torch.cuda.event可能會使統(tǒng)計時間比實際時間短，此時可以使用time.time()代替，標(biāo)準(zhǔn)的with對象書寫如下：

# import time
class Timer:
    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        torch.cuda.synchronize()
        self.elapsed_time = time.time() - self.start_time

然后只需要將上文的with CudaTimer() as timer替換為with Timer() as timer即可。

到此這篇關(guān)于Pytorch精準(zhǔn)記錄函數(shù)運行時間的方法的文章就介紹到這了,更多相關(guān)Pytorch記錄函數(shù)運行時間內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: