来源:DeepHub IMBA本文约1300字,建议阅读5分钟本文将通过一个pytorch GPU 调试的示例来说明如何使用它们。装饰器是 python 上下文管理器的特定实现。本片文章将通过一个pytorch GPU 调试的示例来说明如何使用它们。虽然它可能不适用于所有情况,但我它们却是非常有用。 调试内存泄漏问题有很多方法可以调试内存泄漏。本文将展示一种识别代码中有问题的行的有用方法。该方法可以有助于以简洁的方式找到具体的位置。 逐行手动调试 如果遇到问题,一种经典的且常用的方法是使用调试器逐行检查,比如下面的例子:
它可以工作,但这样的操作光听起来来就很麻烦。我们可以将其封装成一个函数,这样可以在需要的时候调用,这样几乎不需要修改现有的代码,所以就引出了我们要介绍装饰器的功能。 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所做的工作。因为它但更容易开发,也更容易使用:
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定制的,但可以很容易地修改为其他的库:
现在就可以使用了,我们对任何一行(或块)代码使用这个上下文: 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。 |
|