pytorch训练网络时内存不够问题

最近写代码时遇到这个问题，研究了很久，最终还是解决了这个问题，这里放一下解决过程。

由于我这里做的工作并不是完全基于神经网络的，只是利用了反向传播，所以我的前向传播并不同于传统的传播方式，即输入经过网络得到输出，我的输出是经过一个更新公式得到的，即经过此更新公式得到的输出就是神经网络的输出，所以问题也就出在这里。

下图为改动前第十次迭代cpu内存使用

每次迭代时，输入经过更新公式前向传播，内存会泄漏一些，我开始并不知道这是什么原因，以为只是计算所产生的内存占用，这里是用pip install memory_profiler安装的一个监视内存的包，使用方法如下，from memory_profiler import profile调用包，然后在需要查看的函数前加上@profile(precision=4, stream=open('test.log', 'w+'))即可，'test.log'为保存的日志文件名称。

import torch
import torch.nn as nn
from memory_profiler import profile
@profile(precision=4, stream=open('test.log', 'w+'))
def aa(a,b):
    d=1.*torch.mm(a,b)
    return d
a=torch.randn(3,4)
b=torch.randn(4,5)
aa(a,b)
print(a@b)

改动前我的内存占用情况如下图，可以看到每次计算的占用很大，实际上这部分内存并不是完全由计算值所产生，还有一部分是由于梯度的存在所产生，而这部分在梯度更新反向传播的过程中都没有清除，所以产生了梯度保存占用内存的情况，所以只需要清除这部分所产生的内存就行。记事本打开日志文件，看到内存这部分占用了89.7422Mb。

Filename: F:/code/pythonProject/pythonProject_ddl_lpom_learned_net/deep_transform_learning_LPOM_bp_update.py

Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
   134 5726.4102 MiB 5726.4102 MiB           1       @profile(precision=4, stream=open('memory_train.log', 'w+'))
   135                                             def update_z_1_2(self,mu,layer):
   136 5726.4102 MiB   0.0000 MiB           1           z1=self.z_layers[layer+1]
   137 5726.4102 MiB   0.0000 MiB           1           z2=self.z_layers[layer+2]
   138 5726.4102 MiB   0.0000 MiB           1           z0=self.z_layers[layer]
   139 5726.4102 MiB   0.0000 MiB           1           t0=self.transform_layers[layer]
   140 5726.4102 MiB   0.0000 MiB           1           t1=self.transform_layers[layer+1]
   141                                                 #with torch.no_grad():
   142 5816.1523 MiB  89.7422 MiB           1           self.z_layers[layer+1] = leakyrelu(torch.mm(t0, z0) - torch.mm(torch.t(t1), leakyrelu(torch.mm(t1, z1)) - z2))

实际上这是包含了梯度的计算值，而这个梯度值每次迭代都会保存，所以最后会导致内存不够用的情况，由于这个梯度值我不需要，同神经网络中不需要激活值的梯度一样，我们只需要权值的梯度。所以在这句话之前加上了：with torch.no_grad():   结果如下，

Filename: F:/code/pythonProject/pythonProject_ddl_lpom_learned_net/deep_transform_learning_LPOM_bp_update.py

Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
   134 487.0586 MiB 487.0586 MiB           1       @profile(precision=4, stream=open('memory_train.log', 'w+'))
   135                                             def update_z_1_2(self,mu,layer):
   136 487.0586 MiB   0.0000 MiB           1           z1=self.z_layers[layer+1]
   137 487.0586 MiB   0.0000 MiB           1           z2=self.z_layers[layer+2]
   138 487.0586 MiB   0.0000 MiB           1           z0=self.z_layers[layer]
   139 487.0586 MiB   0.0000 MiB           1           t0=self.transform_layers[layer]
   140 487.0586 MiB   0.0000 MiB           1           t1=self.transform_layers[layer+1]
   141                                                 #z1 = leakyrelu(torch.mm(t0,z0) - (mu / mu) * torch.mm(torch.t(t1), leakyrelu(torch.mm(t1,z1)) - z2))
   142 487.0586 MiB   0.0000 MiB           1           with torch.no_grad():
   143 516.9688 MiB  29.9102 MiB           1               self.z_layers[layer+1] = leakyrelu(torch.mm(t0, z0) - torch.mm(torch.t(t1), leakyrelu(torch.mm(t1, z1)) - z2))

可以看到内存的使用减少到29.9102Mb，减少的部分应该是保存的梯度值，这部分我也没有深究。由于每次迭代后这个部分的内存会释放，所以我的代码这个部分就没有出现内存泄漏导致的最终内存不够的情况了。最后展示一下第十次迭代时内存使用情况。最后的迭代结束也是没有出现内存不够的情况，同时因为少计算一部分梯度，迭代一次用的时间也减少了一半。