调试GPU程序错误《CUDA out of memory》

问题描述

我们在跑深度学习程序的时候，经常会遇到CUDA out of memory，一些简单的方案比如

• 缩小Batch Size
• 及时的把不用的变量移动到CPU上
• 执行empty_cache()
• 避免GPU变量累积，例如涉及到GPU变量的跨轮次的加操作和append操作

如果上述的方法依然没有解决问题，可以通过分析GPU上的变量来进行相关的调试：

def debug_memory():
    import collections, gc, resource, torch
    print('maxrss = {}'.format(
        resource.getrusage(resource.RUSAGE_SELF).ru_maxrss))
    tensors = collections.Counter((str(o.device), o.dtype, tuple(o.shape))
                                  for o in gc.get_objects()
                                  if torch.is_tensor(o))
    for line in tensors.items():
        print('{}\t{}'.format(*line))

进一步的，可以完善一下代码：

def debug_memory_sum():
    print('maxrss = {}'.format(
        resource.getrusage(resource.RUSAGE_SELF).ru_maxrss))
    tensors = collections.Counter((str(o.device), o.dtype, tuple(o.shape))
                                  for o in gc.get_objects()
                                  if torch.is_tensor(o))
    memory_total = 0
    for line in tensors.items():
        if line[0][0] == 'cuda:0':
            _shape = line[0][2]
            _cnt = line[1]
            if len(_shape) == 2:
                memory_total = memory_total + _cnt * (_shape[0] * _shape[1])
            elif len(_shape) == 1:
                memory_total = memory_total + _cnt * _shape[0]
            else:
                pass
    print(memory_total)

需要注意的是resource库只能用在linux系统上，然后如果你的算法任务涉及到三维（如图像、视频等）以上，注意修改并添加len(_shape)的判断条件