pytorch 减小显存消耗，优化显存使用，避免out of memory

pytorch 减小显存消耗，优化显存使用，避免out of memory

本文是整理了大神的两篇博客：

如何计算模型以及中间变量的显存占用大小：

https://oldpan.me/archives/how-to-calculate-gpu-memory

如何在Pytorch中精细化利用显存：

https://oldpan.me/archives/how-to-use-memory-pytorch

还有知乎中大神的解答：

https://zhuanlan.zhihu.com/p/31558973

ppt

https://www.zhihu.com/question/67209417

在说之前先推荐一个实时监控内存显存使用的小工具：

sudo apt-get install htop

监控内存（-d为更新频率，下为每0.1s更新一次）：

htop -d=0.1

监控显存（-n为更新频率，下为每0.1s更新一次）：

watch -n 0.1 nvidia-smi

1.问题陈述：

torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58

令人窒息的显存溢出，有时是沉默式gg，不动声色的就溢出没有了。。原因是：

显存装不下模型权重+中间变量

优化方法：及时清空中间变量，优化代码，减少batch

2.显存消耗计算方法：

先看看我们使用的pytorch数据格式：

平时训练中使用的多是float32 和 int32。

32位的单精度浮点型占用空间为4B，

那么一个batch在网络开始比如说是16×3×224×224，那么所占用的显存也就是16×3×224×224×4B = 9.1875MB

到了网络后期比如说是16×512×14*14，所占用的显存也就是16×512×14×14×4B = 6.125MB

即使是256的batch_size，也就是147MB，整个网络如果是19层，为2.728GB，并没有到咱们至少8G的显存。

显存消耗的幕后黑手其实是神经网络中的中间变量以及使用optimizer算法时产生的巨量的中间参数。

显存占用 = 模型参数 + 计算产生的中间变量

以VGG16为例：

原文博主注意到上图中在计算的时候默认的数据格式是8-bit而不是32-bit，所以最后的结果要乘上一个4，即552mb。

其实只要一计算，就可以知道当batch_size是256时，中间变量所产生的参数量是有多庞大。。。

反向传播时，中间变量+原来保存的中间变量，存储量会翻倍。

而且有些适用于移动端的网络mobilenet等，计算量是变少了，但对显存占用变大了，原因就是中间参数存储增加了。

3.代码计算显存占用：

计算模型权重及中间变量占用大小：

1. # 模型显存占用监测函数

2. # model：输入的模型

3. # input：实际中需要输入的Tensor变量

4. # type_size 默认为 4 默认类型为 float32

5.  

6. def modelsize(model, input, type_size=4):

7. para = sum([np.prod(list(p.size())) for p in model.parameters()])

8. print('Model {} : params: {:4f}M'.format(model._get_name(), para * type_size / 1000 / 1000))

9.  

10. input_ = input.clone()

11. input_.requires_grad_(requires_grad=False)

12.  

13. mods = list(model.modules())

14. out_sizes = []

15.  

16. for i in range(1, len(mods)):

17. m = mods[i]

18. if isinstance(m, nn.ReLU):

19. if m.inplace:

20. continue

21. out = m(input_)

22. out_sizes.append(np.array(out.size()))

23. input_ = out

24.  

25. total_nums = 0

26. for i in range(len(out_sizes)):

27. s = out_sizes[i]

28. nums = np.prod(np.array(s))

29. total_nums += nums

30.  
31.  

32. print('Model {} : intermedite variables: {:3f} M (without backward)'

33. .format(model._get_name(), total_nums * type_size / 1000 / 1000))

34. print('Model {} : intermedite variables: {:3f} M (with backward)'

35. .format(model._get_name(), total_nums * type_size*2 / 1000 / 1000))

实际消耗会大一些，因为有框架消耗。

4.其他方法：

a. inplace替换：

我们都知道激活函数Relu()有一个默认参数inplace，默认设置为False，当设置为True时，我们在通过relu()计算时的得到的新值不会占用新的空间而是直接覆盖原来的值，这也就是为什么当inplace参数设置为True时可以节省一部分内存的缘故。

b. 用del一遍计算一边清除中间变量。

c. 用checkpoint牺牲计算速度：

在Pytorch-0.4.0出来了一个新的功能，可以将一个计算过程分成两半，也就是如果一个模型需要占用的显存太大了，我们就可以先计算一半，保存后一半需要的中间结果，然后再计算后一半。

也就是说，新的checkpoint允许我们只存储反向传播所需要的部分内容。如果当中缺少一个输出(为了节省内存而导致的)，checkpoint将会从最近的检查点重新计算中间输出，以便减少内存使用(当然计算时间增加了)：

1. # 首先设置输入的input=>requires_grad=True

2. # 如果不设置可能会导致得到的gradient为0

3.  

4. input = torch.rand(1, 10, requires_grad=True)

5. layers = [nn.Linear(10, 10) for _ in range(1000)]

6.  

7. # 定义要计算的层函数，可以看到我们定义了两个

8. # 一个计算前500个层，另一个计算后500个层

9.  

10. def run_first_half(*args):

11. x = args[0]

12. for layer in layers[:500]:

13. x = layer(x)

14. return x

15.  

16. def run_second_half(*args):

17. x = args[0]

18. for layer in layers[500:-1]:

19. x = layer(x)

20. return x

21.  

22. # 我们引入新加的checkpoint

23. from torch.utils.checkpoint import checkpoint

24.  

25. x = checkpoint(run_first_half, input)

26. x = checkpoint(run_second_half, x)

27. # 最后一层单独调出来执行

28. x = layers[-1](x)

29. x.sum.backward() # 这样就可以了

对于Sequential-model来说，因为Sequential()中可以包含很多的block，所以官方提供了另一个功能包：

1. input = torch.rand(1, 10, requires_grad=True)

2. layers = [nn.Linear(10, 10) for _ in range(1000)]

3. model = nn.Sequential(*layers)

4.  

5. from torch.utils.checkpoint import checkpoint_sequential

6.  

7. # 分成两个部分

8. num_segments = 2

9. x = checkpoint_sequential(model, num_segments, input)

10. x.sum().backward() # 这样就可以了

d.减小batch_size, 避免用全连接，多用下采样。

e. torch.backends.cudnn.benchmark = True 在程序刚开始加这条语句可以提升一点训练速度，没什么额外开销。

f. 因为每次迭代都会引入点临时变量，会导致训练速度越来越慢，基本呈线性增长。开发人员还不清楚原因，但如果周期性的使用torch.cuda.empty_cache()的话就可以解决这个问题。

5.显存跟踪：

开头链接的博主开发了一个库：pynvml（Nvidia的Python环境库和Python的垃圾回收工具）

可以实时地打印我们使用的显存以及哪些Tensor使用了我们的显存

https://github.com/Oldpan/Pytorch-Memory-Utils

1. import datetime

2. import linecache

3. import os

4.  

5. import gc

6. import pynvml

7. import torch

8. import numpy as np

9.  
10.  

11. print_tensor_sizes = True

12. last_tensor_sizes = set()

13. gpu_profile_fn = f'{datetime.datetime.now():%d-%b-%y-%H:%M:%S}-gpu_mem_prof.txt'

14.  

15. # if 'GPU_DEBUG' in os.environ:

16. # print('profiling gpu usage to ', gpu_profile_fn)

17.  

18. lineno = None

19. func_name = None

20. filename = None

21. module_name = None

22.  

23. # fram = inspect.currentframe()

24. # func_name = fram.f_code.co_name

25. # filename = fram.f_globals["__file__"]

26. # ss = os.path.dirname(os.path.abspath(filename))

27. # module_name = fram.f_globals["__name__"]

28.  
29.  

30. def gpu_profile(frame, event):

31. # it is _about to_ execute (!)

32. global last_tensor_sizes

33. global lineno, func_name, filename, module_name

34.  

35. if event == 'line':

36. try:

37. # about _previous_ line (!)

38. if lineno is not None:

39. pynvml.nvmlInit()

40. # handle = pynvml.nvmlDeviceGetHandleByIndex(int(os.environ['GPU_DEBUG']))

41. handle = pynvml.nvmlDeviceGetHandleByIndex(0)

42. meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)

43. line = linecache.getline(filename, lineno)

44. where_str = module_name+' '+func_name+':'+' line '+str(lineno)

45.  

46. with open(gpu_profile_fn, 'a+') as f:

47. f.write(f"At {where_str:<50}"

48. f"Total Used Memory:{meminfo.used/1024**2:<7.1f}Mb\n")

49.  

50. if print_tensor_sizes is True:

51. for tensor in get_tensors():

52. if not hasattr(tensor, 'dbg_alloc_where'):

53. tensor.dbg_alloc_where = where_str

54. new_tensor_sizes = {(type(x), tuple(x.size()), np.prod(np.array(x.size()))*4/1024**2,

55. x.dbg_alloc_where) for x in get_tensors()}

56. for t, s, m, loc in new_tensor_sizes - last_tensor_sizes:

57. f.write(f'+ {loc:<50} {str(s):<20} {str(m)[:4]} M {str(t):<10}\n')

58. for t, s, m, loc in last_tensor_sizes - new_tensor_sizes:

59. f.write(f'- {loc:<50} {str(s):<20} {str(m)[:4]} M {str(t):<10}\n')

60. last_tensor_sizes = new_tensor_sizes

61. pynvml.nvmlShutdown()

62.  

63. # save details about line _to be_ executed

64. lineno = None

65.  

66. func_name = frame.f_code.co_name

67. filename = frame.f_globals["__file__"]

68. if (filename.endswith(".pyc") or

69. filename.endswith(".pyo")):

70. filename = filename[:-1]

71. module_name = frame.f_globals["__name__"]

72. lineno = frame.f_lineno

73.  

74. return gpu_profile

75.  

76. except Exception as e:

77. print('A exception occured: {}'.format(e))

78.  

79. return gpu_profile

80.  
81.  

82. def get_tensors():

83. for obj in gc.get_objects():

84. try:

85. if torch.is_tensor(obj):

86. tensor = obj

87. else:

88. continue

89. if tensor.is_cuda:

90. yield tensor

91. except Exception as e:

92. print('A exception occured: {}'.format(e))

需要注意的是，linecache中的getlines只能读取缓冲过的文件，如果这个文件没有运行过则返回无效值。Python 的垃圾收集机制会在变量没有应引用的时候立马进行回收，但是为什么模型中计算的中间变量在执行结束后还会存在呢。既然都没有引用了为什么还会占用空间？
一种可能的情况是这些引用不在Python代码中，而是在神经网络层的运行中为了backward被保存为gradient，这些引用都在计算图中，我们在程序中是无法看到的。