问题描述
我们在跑深度学习程序的时候,经常会遇到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)
的判断条件