头图

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

装饰器是 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://avoid.overfit.cn/post/40d81e2235d345ed9f25d2221af7cbcf

最后希望这篇小文章能让你了解什么是上下文管理器,如何使用上下文装饰器,以及如何将它们应用于调试pytorch。

作者:MarkTension


deephub
提供专业的CV NLP和数据挖掘知识,更多的干货请关注 公众号 Deephub-IMBA
58 声望
22 粉丝
0 条评论
推荐阅读
NLP / LLMs中的Temperature 是什么?
大型语言模型能够根据给定的上下文或提示生成新文本,由于神经网络等深度学习技术的进步,这些模型越来越受欢迎。可用于控制生成语言模型行为的关键参数之一是Temperature 参数。在本文中,我们将讨论语言生成模...

deephub

封面图
最好用的 python 库合集
🎈 分词 - jieba优秀的中文分词库,依靠中文词库,利用词库确定汉子之间关联的概率,形成分词结果 {代码...} 🎈 词云库 - wordcloud对数据中出现频率较高的 关键词 生成的一幅图像,予以视觉上的突出 {代码...} 🎈 ...

tiny极客11阅读 2.8k评论 2

封面图
数据结构与算法:二分查找
一、常见数据结构简单数据结构(必须理解和掌握)有序数据结构:栈、队列、链表。有序数据结构省空间(储存空间小)无序数据结构:集合、字典、散列表,无序数据结构省时间(读取时间快)复杂数据结构树、 堆图二...

白鲸鱼9阅读 5.2k

滚蛋吧,正则表达式!
你是不是也有这样的操作,比如你需要使用「电子邮箱正则表达式」,首先想到的就是直接百度上搜索一个,然后采用 CV 大法神奇地接入到你的代码中?

良许3阅读 1.4k

超详细的ChatGPT注册教程来了
最近一周,大家都在讨论ChatGPT,一些主流的技术社区更是将ChatGPT吹的神乎其技,那ChatGPT是什么呢?又能给我们带来哪些变化呢?。带着这些问题,我打算先注册并使用 ChatGPT,供想要体验 ChatGPT 的小伙伴们参考。

xiangzhihong3阅读 2.7k评论 3

搭个ChatGPT算法模型,从哪开始?
最近 ChatGPT 很火,火到了各行各业。记得去年更多的还是码农最新体验后拿它搜代码,现在各行各业都进来体验,问它咋理财、怎么写报告和给小孩起名。😂 也因此让小傅哥在头条的一篇关于 ChatGPT 的文章都有了26万...

小傅哥6阅读 1.1k

封面图
程序员适合创业吗?
大家好,我是良许。从去年 12 月开始,我已经在视频号、抖音等主流视频平台上连续更新视频到现在,并得到了不错的评价。每个视频都花了很多时间精力用心制作,欢迎大家关注哦~考虑到有些小伙伴没有看过我的视频,...

良许3阅读 1.2k

58 声望
22 粉丝
宣传栏