裝飾器是 python 上下文管理器的特定實作。本片文章將透過一個pytorch GPU 調試的範例來說明如何使用它們。雖然它可能不適用於所有情況,但我它們卻非常有用。
有很多方法可以偵錯記憶體洩漏。本文將展示一種識別程式碼中有問題的行的有用方法。此方法可以有助於以簡潔的方式找到具體的位置。
如果遇到問題,一種經典的且常用的方法是使用偵錯器逐行檢查,例如下面的例子:
它可以工作,但這樣的操作光聽起來來就很麻煩。我們可以將其封裝成函數,這樣可以在需要的時候調用,這樣幾乎不需要修改現有的程式碼,所以就引出了我們要介紹裝飾器的功能。
裝飾器可以包裝在程式碼的任意部分。這裡我們透過裝飾器來完成檢查是否有額外的張量 ,除此以外我們還需要一個計數器,因為需要在執行之前和之後計算張量的數量。模式如下所示:
def memleak_wrapper(func): def wrap(*args, **kwargs): print("num tensors start is ...") out = func(*args, **kwargs) print("num tensors end is ...") return out return wrap@memleak_wrapper def function_to_debug(x): print(f"put line(s) of code here. Input is {x}") out = x + 10 return outout = function_to_debug(x=1000) print(f"out is {out}") #输入类似这样 #num tensors start is ... #put line(s) of code here. Input is 1000 #num tensors end is ... #outis 1010
要執行這個程式碼,我們就需要將要檢查的程式碼行放入函數 (function_to_debug)。但這不是最好的,因為我們還需要手動插入很多程式碼。另外就是如果程式碼區塊產生的變數多於一個,還需要尋找額外的解決方案來使用這些下游變數。
為了解決上面問題,我們的可以使用上下文管理器來取代函數裝飾器。上下文管理器最廣泛使用的範例是使用 with 語句實例化上下文。以前最常見的就是:
with open("file") as f: …
使用Python的contextlib函式庫,Python使用者可以輕鬆地自行建立上下文管理器。所以本文中我們將使用ContextDecorator完成在上面嘗試使用decorator所做的工作。因為它但更容易開發,也更容易使用:
from contextlib import ContextDecorator class check_memory_leak_context(ContextDecorator): def __enter__(self): print('Starting') return self def __exit__(self, *exc): print('Finishing') return False
ContextDecorator 有2 個方法:enter() 和 exit() ,當我們進入或退出上下文時會呼叫它們。 __exit__ 中的 *exc 參數代表任何傳入的例外。
現在我們來用它來解決上面說的問題。
因為要計算張量的總數,所以我們將計算過程封裝成一個函數get_n_tensors() ,這樣可以在上下文開始和結束時來計算張數量:
class check_memory_leak_context(ContextDecorator): def __enter__(self): self.start = get_n_tensors() return self def __exit__(self, *exc): self.end = get_n_tensors() increase = self.end — self.start if increase > 0: print(f”num tensors increased with" f"{self.end — self.start} !”) else: print(”no added tensors”) return False
如果有增加,則將其列印到控制台。
get_n_tensor()使用垃圾收集器(gc),是為pytorch定制的,但可以輕鬆修改為其他的庫:
import gc def get_n_tensors(): tensors= [] for obj in gc.get_objects(): try: if (torch.is_tensor(obj) or (hasattr(obj, ‘data’) and torch.is_tensor(obj.data))): tensors.append(obj) except: pass return len(tensors)
現在就可以使用了,我們對任何一行(或區塊)程式碼使用這個上下文:
x = arbitrary_operation(x) ... with check_memory_leak_context(): y = x[0].permute(1, 2, 0).cpu().detach().numpy() x = some_harmless_operation() ... x = another_arbitrary_operation(x)
如果上下文修飾器包裝的行內建立了一個新的張量,它就會列印出來。
這是一個非常好的程式碼片段,你可以在開發過程中把它放在一個單獨的檔案中,下面是本文的完整程式碼:
https://gist.github.com/MarkTension/4783697ebd5212ba500cdd829b364338
最後希望這篇小文章能讓你了解什麼是上下文管理器,如何使用上下文裝飾器,以及如何將它們應用於調試pytorch 。
以上是使用上下文裝飾器調試Pytorch的記憶體洩漏問題的詳細內容。更多資訊請關注PHP中文網其他相關文章!