使用上下文装饰器调试Pytorch的内存泄漏问题

2023-05-14 20:05:35 上下文 泄漏 调试

装饰器是 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。

以上就是使用上下文装饰器调试Pytorch的内存泄漏问题的详细内容，更多请关注其它相关文章！

相关文章