快捷導(dǎo)航

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

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

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

0. 引言

參考Pytorch官方文檔對(duì)CUDA的描述，GPU的運(yùn)算是異步執(zhí)行的。一般來(lái)說(shuō)，異步計(jì)算的效果對(duì)于調(diào)用者來(lái)說(shuō)是不可見的，因?yàn)?/p>

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

異步計(jì)算的后果是，沒(méi)有同步的時(shí)間測(cè)量是不準(zhǔn)確的。

1. 解決方案

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

# 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()中間以測(cè)量時(shí)間（單位毫秒）。

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

先書寫一個(gè)自定義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語(yǔ)句返回：

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

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

2. 補(bǔ)充

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

# 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ù)運(yùn)行時(shí)間的方法的文章就介紹到這了,更多相關(guān)Pytorch記錄函數(shù)運(yùn)行時(shí)間內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: