装饰器是python上下文管理器的特定实现。本文将通过一个pytorchGPU调试的例子来说明如何使用它们。虽然它可能并非在所有情况下都有效,但我认为它们非常有用。调试内存泄漏问题调试内存泄漏的方法有很多种。本文将展示一种用于识别代码中有问题的行的有用方法。这种方法可以帮助以简洁的方式找到特定位置。逐行手动调试如果遇到问题,经典且常用的方法是使用调试器逐行检查,例如下面的例子:在搜索引擎中找到有关pytorch中如何计算张量总数的代码片段,如:tensor-counter-snippet在代码中设置断点使用tensor-counter-snippet获取张量总数使用调试器执行下一步重新运行tensor-counter-snippet并检查张量计数是否增加重复上面的步骤就可以了,但是这样的操作听起来比较麻烦。我们可以把它封装成一个函数,需要的时候调用,这样几乎不需要修改现有的代码,所以就引出了装饰器的功能。Python装饰器装饰器可以包裹代码的任何部分。这里我们通过装饰器检查是否有多余的张量,我们还需要一个计数器,因为执行前后需要计算张量的数量。模式如下所示:defmemleak_wrapper(func):defwrap(*args,**kwargs):print("numtensorsstartis...")out=func(*args,**kwargs)print("num张量结束是...")returnoutreturnwrap@memleak_wrapperdeffunction_to_debug(x):print(f"putline(s)ofcodehere.Inputis{x}")out=x+10returnoutout=function_to_debug(x=1000)print(f"outis{out}")#Inputlikethis#numtensorsstartis...#putline(s)ofcodehere.Inputis1000#numtensorsendis...#outis1010要运行此代码,我们需要将要检查的代码行放入函数(function_to_debug)中。但这并不是最好的,因为我们还需要手动插入很多代码。此外,如果代码块生成多个变量,则需要寻找额外的解决方案来使用这些下游变量。上下文装饰器为了解决上述问题,我们可以使用上下文管理器来代替函数装饰器。上下文管理器使用最广泛的示例是使用with语句实例化上下文。最常见的曾经是:withopen("file")asf:...使用Python的contextlib库,Python用户可以轻松创建自己的上下文管理器。因此,在本文中,我们将使用ContextDecorator来完成我们在上面尝试使用装饰器所做的事情。因为更容易开发和使用:fromcontextlibimportContextDecoratorclasscheck_memory_leak_context(ContextDecorator):def__enter__(self):print('Starting')returnselfdef__exit__(self,*exc):print('Finishing')returnFalseContextDecorator有2个方法:enter()和exit(),当我们进入或退出上下文时调用它们。__exit__中的*exc参数表示任何传入的异常。现在让我们用它来解决上面提到的问题。使用ContextDecorator查找内存泄漏因为需要计算张量的总数,所以我们将计算过程封装成一个函数get_n_tensors(),可以计算上下文首尾张量的个数:classcheck_memory_leak_context(ContextDecorator):def__enter__(self):self.start=get_n_tensors()returnselfdef__exit__(self,*exc):self.end=get_n_tensors()increase=self.end—self.startifincrease>0:print(f”numtensorsincreasedwith"\f"{self.end—self.start}!")else:print("noaddedtensors")returnFalse如果有加法,打印到控制台。get_n_tensor()使用垃圾收集器(gc),是为pytorch定制的,但可以很容易地为其他库修改:importgcdefget_n_tensors():tensors=[]forobjingc.get_objects():.is_tensor(obj)or(hasattr(obj,'data')andtorch.is_tensor(obj.data))):tensors.append(obj)except:passreturnlen(tensors)现在有效,我们将此上下文用于任何代码行(或代码块):x=arbitrary_operation(x)...withcheck_memory_leak_context():y=x[0].permute(1,2,0).cpu().detach().numpy()x=some_harmless_operation()...x=another_arbitrary_operation(x)如果在上下文装饰器包裹的行内创建了新的张量,则将打印。总结这是一个非常好的代码片段,开发时可以放在单独的文件中,这里是本文的完整代码:https://avoid.overfit.cn/post/40d81e2235d345ed9f25d2221af7cbcf最后希望这篇小文章能够让你了解上下文管理器是什么,如何使用上下文装饰器,以及如何将它们应用于调试pytorch。添加一名作者
