來源:DeepHub IMBA本文約1300字,建議閱讀5分鐘本文将通過一個pytorch GPU 調試的示例來說明如何使用它們。
裝飾器是 python 上下文管理器的特定實作。本片文章将通過一個pytorch GPU 調試的示例來說明如何使用它們。雖然它可能不适用于所有情況,但我它們卻是非常有用。
調試記憶體洩漏問題
有很多方法可以調試記憶體洩漏。本文将展示一種識别代碼中有問題的行的有用方法。該方法可以有助于以簡潔的方式找到具體的位置。
逐行手動調試
如果遇到問題,一種經典的且常用的方法是使用調試器逐行檢查,比如下面的例子:
- 在搜尋引擎查找有關如何計算 pytorch 中所有張量總數的代碼片段,比如:tensor-counter-snippet
- 在代碼中設定斷點
- 使用tensor-counter-snippet來獲得張量的總數統計
- 使用調試器執行下一步操作
- 重新運作 tensor-counter-snippet,并檢查張量計數是否增加
- 重複上面的步驟
它可以工作,但這樣的操作光聽起來來就很麻煩。我們可以将其封裝成一個函數,這樣可以在需要的時候調用,這樣幾乎不需要修改現有的代碼,是以就引出了我們要介紹裝飾器的功能。
Python 裝飾器
裝飾器可以包裝在代碼的任意部分。這裡我們通過裝飾器來完成檢查是否有額外的張量 ,除此以外我們還需要一個計數器,因為需要在執行之前和之後計算張量的數量。該模式如下所示:
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 參數代表任何傳入的異常。
現在我們來使用它解決上面說的問題。
使用 ContextDecorator 找出記憶體洩漏
因為要計算張量的總數,是以我們将計算過程封裝成一個函數 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。